{
  "cells": [
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "id": "RMivPb5DZwpL"
      },
      "outputs": [],
      "source": [
        "#@title Imports\n",
        "\n",
        "# Licensed under the Apache License, Version 2.0\n",
        "\n",
        "import json\n",
        "import io\n",
        "import random\n",
        "import time\n",
        "import numpy as np\n",
        "import statistics\n",
        "import tensorflow.compat.v1 as tf\n",
        "\n",
        "tf.enable_eager_execution()\n",
        "assert tf.executing_eagerly()\n",
        "\n",
        "print(tf.version.VERSION)\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "id": "P_0tI6aNlvEF"
      },
      "outputs": [],
      "source": [
        "#@title General dataset utility functions\n",
        "\n",
        "DATASET_ADDITION_ALIGNED = \"addition_aligned\"\n",
        "DATASET_ADDITION_NEGATIVES = \"addition_negatives\"\n",
        "DATASET_REVERSING = \"reversing-new\"\n",
        "DATASET_DUPLICATION = \"duplication\"\n",
        "DATASET_CARTESIAN = \"cartesian\"\n",
        "DATASET_INTERSECTION_BOOLEAN = \"intersection_boolean-new\"\n",
        "DATASET_SCAN_LENGTH = \"scan_length\"\n",
        "DATASET_SCAN_ADD_JUMP = \"scan_add_jump\"\n",
        "DATASET_PCFG_PRODUCTIVITY = \"pcfg_productivity\"\n",
        "DATASET_PCFG_SYSTEMATICITY = \"pcfg_systematicity\"\n",
        "DATASET_COGS_FULL = \"cogs_full\"\n",
        "DATASET_CFQ_MCD1 = \"cfq_mcd1\"\n",
        "DATASET_CFQ_MCD1_INTERMEDIATE = \"cfq_mcd1_intermediate\"\n",
        "\n",
        "\n",
        "# General tokens (common to all datasets):\n",
        "PAD_TOKEN = \"[PAD]\"\n",
        "SEP_TOKEN = \"[SEP]\"\n",
        "END_TOKEN = \"[END]\"\n",
        "START_TOKEN = \"[START]\"\n",
        "END_ITERATION_TOKEN = \"[ENDITER]\"\n",
        "\n",
        "PAD_TOKEN_IDX = 0\n",
        "SEP_TOKEN_IDX = 1\n",
        "END_TOKEN_IDX = 2\n",
        "START_TOKEN_IDX = 3\n",
        "END_ITERATION_TOKEN_IDX = 4\n",
        "\n",
        "\n",
        "def max_length(tensor):\n",
        "  return max(len(t) for t in tensor)\n",
        "\n",
        "\n",
        "def decode(seq, vocab):\n",
        "  \"\"\"Decodes a sequence of token IDs to a string.\"\"\"\n",
        "  out = \"\"\n",
        "  for tok in seq:\n",
        "    if tok == 0:\n",
        "      return out\n",
        "    out += vocab[tok] + \" \"\n",
        "  return out\n",
        "\n",
        "\n",
        "def create_dataset_tensors(examples_in_raw, examples_out_raw, max_len_inp, max_len_targ, vocab_to_int):\n",
        "  \"\"\"Translates a dataset to padded token ID tensors.\"\"\"\n",
        "  in_tensor = []\n",
        "  for example in examples_in_raw:\n",
        "    tensor = [vocab_to_int[x] for x in example]\n",
        "    in_tensor.append(tensor)\n",
        "  out_tensor = []\n",
        "  for example in examples_out_raw:\n",
        "    tensor = [vocab_to_int[x] for x in example]\n",
        "    out_tensor.append(tensor)\n",
        "  in_tensor = tf.keras.preprocessing.sequence.pad_sequences(in_tensor, padding='post', maxlen=max_len_inp)\n",
        "  out_tensor = tf.keras.preprocessing.sequence.pad_sequences(out_tensor, padding='post', maxlen=max_len_targ)\n",
        "\n",
        "  print(in_tensor.shape)\n",
        "  print(out_tensor.shape)\n",
        "\n",
        "  return in_tensor, out_tensor\n",
        "\n",
        "\n",
        "def prepare_tf_dataset_tensors(vocab, vocab_to_int, input_tensor_train, target_tensor_train,\n",
        "                               input_tensor_val_list, target_tensor_val_list,\n",
        "                               batch_size):\n",
        "  \"\"\"Converts dataset tensors to TensorFlow datasets.\"\"\"\n",
        "\n",
        "  BUFFER_SIZE = len(input_tensor_train)\n",
        "\n",
        "  # Print some info for visual inspection:\n",
        "  print(f\"Training set: {len(input_tensor_train)} examples, of length {len(input_tensor_train[0])} -\u003e {len(target_tensor_train[0])}\")\n",
        "  for i in range(len(input_tensor_val_list)):\n",
        "    input_tensor_test = input_tensor_val_list[i]\n",
        "    target_tensor_test = target_tensor_val_list[i]\n",
        "    print(f\"Test set: {len(input_tensor_test)} examples, of length {len(input_tensor_test[0])} -\u003e {len(target_tensor_test[0])}\")\n",
        "  print(f\"Example Input 1: {decode(input_tensor_train[0], vocab)}\")\n",
        "  print(f\"Example Output 1: {decode(target_tensor_train[0], vocab)}\")\n",
        "  print(f\"Example Input 2: {decode(input_tensor_train[-1], vocab)}\")\n",
        "  print(f\"Example Output 2: {decode(target_tensor_train[-1], vocab)}\")\n",
        "  print(f\"Vocab size: {len(vocab)}\")\n",
        "  print(f\"Vocab: {vocab}\")\n",
        "\n",
        "  dataset_train = tf.data.Dataset.from_tensor_slices((input_tensor_train, target_tensor_train)).shuffle(BUFFER_SIZE)\n",
        "  dataset_train = dataset_train.batch(batch_size, drop_remainder=True)\n",
        "\n",
        "  dataset_val_list = []\n",
        "\n",
        "  for i in range(len(input_tensor_val_list)):\n",
        "    dataset_val = tf.data.Dataset.from_tensor_slices((input_tensor_val_list[i],\n",
        "                                                      target_tensor_val_list[i])).shuffle(BUFFER_SIZE)\n",
        "    dataset_val = dataset_val.batch(batch_size, drop_remainder=True)\n",
        "    dataset_val_list.append(dataset_val)\n",
        "  return (dataset_train, dataset_val_list)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "id": "jpOn10oJLShM"
      },
      "outputs": [],
      "source": [
        "#@title \"Addition\" dataset generation\n",
        "\n",
        "def create_addition_dataset(trainsize, testsize, vocab, vocab_to_int,\n",
        "                            reversedigits=True, leftpadding=12,\n",
        "                            addAlignmentTokens=False, negativeProbability=0.0):\n",
        "  \"\"\"Generates the addition dataset.\"\"\"\n",
        "\n",
        "  max_number_length = 12\n",
        "\n",
        "  def add_alignment(digits):\n",
        "    digits2 = []\n",
        "    pos = len(digits)\n",
        "    for digit in digits:\n",
        "      digits2.append(digit)\n",
        "      digits2.append(\"P\" + str(pos))\n",
        "      pos -= 1\n",
        "    return digits2\n",
        "\n",
        "  def create_example(minlen, maxlen, leftpadding=0, reversedigits=False,\n",
        "                     addAlignmentTokens=False, negativeProbability=0.0):\n",
        "    numbers = [0, 1, 2, 3, 4, 5, 6, 7, 8, 9]\n",
        "    digits1 = []\n",
        "    digits2 = []\n",
        "    digitssum = []\n",
        "    int1 = 0\n",
        "    int2 = 0\n",
        "    l1 = random.randint(minlen, maxlen)\n",
        "    l2 = random.randint(minlen, maxlen)\n",
        "    for i in range(l1):\n",
        "      n = random.choice(numbers)\n",
        "      int1 = int1*10 + n\n",
        "      digits1.append(str(n))\n",
        "    for i in range(l2):\n",
        "      n = random.choice(numbers)\n",
        "      int2 = int2*10 + n\n",
        "      digits2.append(str(n))\n",
        "    if random.random() \u003c negativeProbability:\n",
        "      int1 = -int1\n",
        "      digits1 = [\"-\"] + digits1\n",
        "    if random.random() \u003c negativeProbability:\n",
        "      int2 = -int2\n",
        "      digits2 = [\"-\"] + digits2\n",
        "\n",
        "    sum = int1 + int2\n",
        "    negatedSum = False\n",
        "    sumtmp = sum\n",
        "    if sumtmp \u003c 0:\n",
        "      negatedSum = True\n",
        "      sumtmp = -sumtmp\n",
        "    while sumtmp \u003e 0:\n",
        "      digitssum = [str(sumtmp % 10)] + digitssum\n",
        "      sumtmp //= 10\n",
        "    if negatedSum:\n",
        "      digitssum = [\"-\"] + digitssum\n",
        "\n",
        "    leftpaddingtoken = \"0\"\n",
        "    if negativeProbability \u003e 0:\n",
        "      leftpaddingtoken = \"#\"\n",
        "\n",
        "    while len(digits1) \u003c leftpadding:\n",
        "      digits1 = [leftpaddingtoken] + digits1\n",
        "    while len(digits2) \u003c leftpadding:\n",
        "      digits2 = [leftpaddingtoken] + digits2\n",
        "    while len(digitssum) \u003c leftpadding:\n",
        "      digitssum = [leftpaddingtoken] + digitssum\n",
        "\n",
        "    if addAlignmentTokens:\n",
        "      digits1 = add_alignment(digits1)\n",
        "      digits2 = add_alignment(digits2)\n",
        "\n",
        "    if reversedigits:\n",
        "      digits1.reverse()\n",
        "      digits2.reverse()\n",
        "      digitssum.reverse()\n",
        "    example_in = digits1 + [SEP_TOKEN] + digits2 + [END_TOKEN]\n",
        "    example_out = [START_TOKEN] + digitssum + [END_TOKEN]\n",
        "\n",
        "    return example_in, example_out\n",
        "\n",
        "\n",
        "  def create_examples(n, minlen, maxlen, leftpadding=0, reversedigits=False,\n",
        "                      addAlignmentTokens=False,negativeProbability=0.0):\n",
        "    examples_in = []\n",
        "    examples_out = []\n",
        "    for i in range(n):\n",
        "      ein, eout = create_example(minlen, maxlen, leftpadding=leftpadding,\n",
        "                                 reversedigits=reversedigits,\n",
        "                                 addAlignmentTokens=addAlignmentTokens,\n",
        "                                 negativeProbability=negativeProbability)\n",
        "      examples_in.append(ein)\n",
        "      examples_out.append(eout)\n",
        "    return examples_in, examples_out\n",
        "\n",
        "  vocab = [PAD_TOKEN, SEP_TOKEN, END_TOKEN, START_TOKEN,\n",
        "           \"0\", \"1\", \"2\", \"3\", \"4\", \"5\", \"6\", \"7\", \"8\", \"9\"]\n",
        "  vocab_to_int = {PAD_TOKEN:0, SEP_TOKEN:SEP_TOKEN_IDX,\n",
        "                  END_TOKEN:END_TOKEN_IDX,\n",
        "                  START_TOKEN:START_TOKEN_IDX,\n",
        "                  \"0\":4, \"1\":5, \"2\":6, \"3\":7, \"4\":8,\n",
        "                  \"5\":9, \"6\":10, \"7\":11, \"8\":12, \"9\":13}\n",
        "  if addAlignmentTokens:\n",
        "    for i in range(max_number_length):\n",
        "      vocab_to_int[\"P\" + str(i)] = len(vocab)\n",
        "      vocab.append(\"P\" + str(i));\n",
        "  if negativeProbability \u003e 0:\n",
        "    vocab_to_int[\"-\"] = len(vocab)\n",
        "    vocab.append(\"-\")\n",
        "    vocab_to_int[\"#\"] = len(vocab)\n",
        "    vocab.append(\"#\")\n",
        "\n",
        "  train_examples_in_raw, train_examples_out_raw = create_examples(\n",
        "      trainsize, 1, 8, leftpadding=leftpadding, reversedigits=reversedigits,\n",
        "      addAlignmentTokens=addAlignmentTokens,\n",
        "      negativeProbability=negativeProbability)\n",
        "  test_easy_examples_in_raw, test_easy_examples_out_raw = create_examples(\n",
        "      testsize, 5, 6, leftpadding=leftpadding, reversedigits=reversedigits,\n",
        "      addAlignmentTokens=addAlignmentTokens,\n",
        "      negativeProbability=negativeProbability)\n",
        "  test_examples_in_raw, test_examples_out_raw = create_examples(\n",
        "      testsize, 6, 8, leftpadding=leftpadding, reversedigits=reversedigits,\n",
        "      addAlignmentTokens=addAlignmentTokens,\n",
        "      negativeProbability=negativeProbability)\n",
        "  test_hard_examples_in_raw, test_hard_examples_out_raw = create_examples(\n",
        "      testsize, 9, 10, leftpadding=leftpadding, reversedigits=reversedigits,\n",
        "      addAlignmentTokens=addAlignmentTokens,\n",
        "      negativeProbability=negativeProbability)\n",
        "\n",
        "  max_len_inp = max(max_length(train_examples_in_raw),\n",
        "                    max_length(test_easy_examples_in_raw),\n",
        "                    max_length(test_hard_examples_in_raw))\n",
        "  max_len_targ = max(max_length(train_examples_out_raw),\n",
        "                     max_length(test_easy_examples_out_raw),\n",
        "                     max_length(test_hard_examples_out_raw))\n",
        "\n",
        "  input_tensor_train, target_tensor_train = create_dataset_tensors(\n",
        "      train_examples_in_raw, train_examples_out_raw, max_len_inp, max_len_targ,\n",
        "      vocab_to_int)\n",
        "  input_tensor_val0, target_tensor_val0 = create_dataset_tensors(\n",
        "      train_examples_in_raw[0:testsize], train_examples_out_raw[0:testsize],\n",
        "      max_len_inp, max_len_targ, vocab_to_int)\n",
        "  input_tensor_val1, target_tensor_val1 = create_dataset_tensors(\n",
        "      test_easy_examples_in_raw, test_easy_examples_out_raw, max_len_inp,\n",
        "      max_len_targ, vocab_to_int)\n",
        "  input_tensor_val2, target_tensor_val2 = create_dataset_tensors(\n",
        "      test_examples_in_raw, test_examples_out_raw, max_len_inp, max_len_targ,\n",
        "      vocab_to_int)\n",
        "  input_tensor_val3, target_tensor_val3 = create_dataset_tensors(\n",
        "      test_hard_examples_in_raw, test_hard_examples_out_raw, max_len_inp,\n",
        "      max_len_targ, vocab_to_int)\n",
        "\n",
        "  return (vocab, vocab_to_int,\n",
        "          input_tensor_train, target_tensor_train,\n",
        "          [input_tensor_val0, input_tensor_val1,\n",
        "           input_tensor_val2, input_tensor_val3],\n",
        "          [target_tensor_val0,target_tensor_val1,\n",
        "           target_tensor_val2, target_tensor_val3])\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "id": "bTXm8AgJcdsY"
      },
      "outputs": [],
      "source": [
        "#@title Algorithmic dataset generation (reversing, duplicating, cartesian, intersection)\n",
        "\n",
        "def create_reversing_dataset(trainsize, testsize, vocab, vocab_to_int,\n",
        "                             trainmindigits=1,trainmaxdigits=16,\n",
        "                             testmindigits=17, testmaxdigits=24):\n",
        "  \"\"\"Generate the reversing dataset.\"\"\"\n",
        "\n",
        "  def create_example(minlen, maxlen):\n",
        "    digits = [\"0\", \"1\", \"2\", \"3\", \"4\", \"5\", \"6\", \"7\", \"8\", \"9\"]\n",
        "    l1 = random.randint(minlen, maxlen)\n",
        "    example_in = []\n",
        "    for i in range(l1):\n",
        "      example_in.append(random.choice(digits))\n",
        "    example_out = example_in[::-1]\n",
        "    example_in.append(END_TOKEN)\n",
        "    example_out = [START_TOKEN] + example_out + [END_TOKEN]\n",
        "    return example_in, example_out\n",
        "\n",
        "\n",
        "  def create_examples(n, minlen, maxlen):\n",
        "    examples_in = []\n",
        "    examples_out = []\n",
        "    for i in range(n):\n",
        "      ein, eout = create_example(minlen, maxlen)\n",
        "      examples_in.append(ein)\n",
        "      examples_out.append(eout)\n",
        "    return examples_in, examples_out\n",
        "\n",
        "  vocab = [PAD_TOKEN, SEP_TOKEN, END_TOKEN, START_TOKEN,\n",
        "           \"0\", \"1\", \"2\", \"3\", \"4\", \"5\", \"6\", \"7\", \"8\", \"9\"]\n",
        "  vocab_to_int = {PAD_TOKEN:0, SEP_TOKEN:SEP_TOKEN_IDX,\n",
        "                  END_TOKEN:END_TOKEN_IDX,\n",
        "                  START_TOKEN:START_TOKEN_IDX,\n",
        "                  \"0\":4, \"1\":5, \"2\":6, \"3\":7, \"4\":8,\n",
        "                  \"5\":9, \"6\":10, \"7\":11, \"8\":12, \"9\":13}\n",
        "\n",
        "  train_examples_in_raw, train_examples_out_raw = create_examples(\n",
        "      trainsize, trainmindigits, trainmaxdigits)\n",
        "  test_easy_examples_in_raw, test_easy_examples_out_raw = create_examples(\n",
        "      testsize, trainmindigits, trainmaxdigits)\n",
        "  test_hard_examples_in_raw, test_hard_examples_out_raw = create_examples(\n",
        "      testsize, testmindigits, testmaxdigits)\n",
        "\n",
        "  max_len_inp = max(max_length(train_examples_in_raw), max_length(\n",
        "      test_easy_examples_in_raw), max_length(test_hard_examples_in_raw))\n",
        "  max_len_targ = max(max_length(train_examples_out_raw), max_length(\n",
        "      test_easy_examples_out_raw), max_length(test_hard_examples_out_raw))\n",
        "\n",
        "  input_tensor_train, target_tensor_train = create_dataset_tensors(\n",
        "      train_examples_in_raw, train_examples_out_raw, max_len_inp, max_len_targ,\n",
        "      vocab_to_int)\n",
        "  input_tensor_val0, target_tensor_val0 = create_dataset_tensors(\n",
        "      train_examples_in_raw[0:testsize], train_examples_out_raw[0:testsize],\n",
        "      max_len_inp, max_len_targ, vocab_to_int)\n",
        "  input_tensor_val1, target_tensor_val1 = create_dataset_tensors(\n",
        "      test_easy_examples_in_raw, test_easy_examples_out_raw, max_len_inp,\n",
        "      max_len_targ, vocab_to_int)\n",
        "  input_tensor_val2, target_tensor_val2 = create_dataset_tensors(\n",
        "      test_hard_examples_in_raw, test_hard_examples_out_raw, max_len_inp,\n",
        "      max_len_targ, vocab_to_int)\n",
        "\n",
        "  return (vocab, vocab_to_int,\n",
        "          input_tensor_train, target_tensor_train,\n",
        "          [input_tensor_val0, input_tensor_val1,\n",
        "           input_tensor_val2],\n",
        "          [target_tensor_val0,target_tensor_val1,\n",
        "           target_tensor_val2])\n",
        "\n",
        "\n",
        "# variableDuplication = 0: just duplicate input (used in paper)\n",
        "# variableDuplication = 1: randomly duplicate 1, 2, 3 or 4 times\n",
        "#                          (with an integer)\n",
        "# variableDuplication = 2: same as 1, but instead of an integer, it's indicated\n",
        "#                          with \"1\", \"1,1\", \"1,1,1\" or \"1,1,1,1\"\n",
        "# variableDuplication = 3: same as 2, but test asks for more duplications than\n",
        "#                          train (this is almost the same as cartesian product).\n",
        "def create_duplicating_dataset(trainsize, testsize, vocab, vocab_to_int,\n",
        "                               trainmindigits=1, trainmaxdigits=16,\n",
        "                               testmindigits=17, testmaxdigits=24,\n",
        "                               variableduplication=0):\n",
        "  \"\"\"Creates the duplicating dataset.\"\"\"\n",
        "\n",
        "  def create_example(minlen, maxlen, minduplications=1, maxduplications=4):\n",
        "    digits = [\"0\", \"1\", \"2\", \"3\", \"4\", \"5\", \"6\", \"7\", \"8\", \"9\"]\n",
        "    l1 = random.randint(minlen, maxlen)\n",
        "    example_in_prefix = []\n",
        "    example_in = []\n",
        "    n_duplications = 2\n",
        "    if variableduplication == 1:\n",
        "      n_duplications = random.choice([1,2,3,4])\n",
        "      example_in_prefix = [digits[n_duplications], SEP_TOKEN]\n",
        "    elif variableduplication == 2 or variableduplication == 3:\n",
        "      n_duplications = random.randint(minduplications, maxduplications)\n",
        "      for i in range(n_duplications):\n",
        "        example_in_prefix.append(digits[1])\n",
        "      example_in_prefix.append(SEP_TOKEN)\n",
        "\n",
        "    for i in range(l1):\n",
        "      example_in.append(random.choice(digits))\n",
        "    example_out = []\n",
        "    for i in range(n_duplications):\n",
        "      example_out += example_in\n",
        "    example_in = example_in_prefix + example_in + [END_TOKEN]\n",
        "    example_out = [START_TOKEN] + example_out + [END_TOKEN]\n",
        "    return example_in, example_out\n",
        "\n",
        "\n",
        "  def create_examples(n, minlen, maxlen, minduplications=1, maxduplications=4):\n",
        "    examples_in = []\n",
        "    examples_out = []\n",
        "    for i in range(n):\n",
        "      ein, eout = create_example(minlen, maxlen,\n",
        "                                 minduplications=minduplications,\n",
        "                                 maxduplications=maxduplications)\n",
        "      examples_in.append(ein)\n",
        "      examples_out.append(eout)\n",
        "    return examples_in, examples_out\n",
        "\n",
        "  vocab = [PAD_TOKEN, SEP_TOKEN, END_TOKEN, START_TOKEN,\n",
        "           \"0\", \"1\", \"2\", \"3\", \"4\", \"5\", \"6\", \"7\", \"8\", \"9\"]\n",
        "  vocab_to_int = {PAD_TOKEN:0, SEP_TOKEN:SEP_TOKEN_IDX,\n",
        "                  END_TOKEN:END_TOKEN_IDX,\n",
        "                  START_TOKEN:START_TOKEN_IDX,\n",
        "                  \"0\":4, \"1\":5, \"2\":6, \"3\":7, \"4\":8,\n",
        "                  \"5\":9, \"6\":10, \"7\":11, \"8\":12, \"9\":13}\n",
        "\n",
        "  if variableduplication == 3:\n",
        "    train_examples_in_raw, train_examples_out_raw = create_examples(\n",
        "        trainsize, trainmindigits, trainmaxdigits)\n",
        "    test_easy_examples_in_raw, test_easy_examples_out_raw = create_examples(\n",
        "        testsize, trainmindigits, trainmaxdigits)\n",
        "    test_hard_examples_in_raw, test_hard_examples_out_raw = create_examples(\n",
        "        testsize, testmindigits, testmaxdigits, minduplications=5,\n",
        "        maxduplications=6)\n",
        "  else:\n",
        "    train_examples_in_raw, train_examples_out_raw = create_examples(\n",
        "        trainsize, trainmindigits, trainmaxdigits)\n",
        "    test_easy_examples_in_raw, test_easy_examples_out_raw = create_examples(\n",
        "        testsize, trainmindigits, trainmaxdigits)\n",
        "    test_hard_examples_in_raw, test_hard_examples_out_raw = create_examples(\n",
        "        testsize, testmindigits, testmaxdigits)\n",
        "\n",
        "  max_len_inp = max(max_length(train_examples_in_raw), max_length(\n",
        "      test_easy_examples_in_raw), max_length(test_hard_examples_in_raw))\n",
        "  max_len_targ = max(max_length(train_examples_out_raw), max_length(\n",
        "      test_easy_examples_out_raw), max_length(test_hard_examples_out_raw))\n",
        "\n",
        "  input_tensor_train, target_tensor_train = create_dataset_tensors(\n",
        "      train_examples_in_raw, train_examples_out_raw, max_len_inp, max_len_targ,\n",
        "      vocab_to_int)\n",
        "  input_tensor_val0, target_tensor_val0 = create_dataset_tensors(\n",
        "      train_examples_in_raw[0:testsize], train_examples_out_raw[0:testsize],\n",
        "      max_len_inp, max_len_targ, vocab_to_int)\n",
        "  input_tensor_val1, target_tensor_val1 = create_dataset_tensors(\n",
        "      test_easy_examples_in_raw, test_easy_examples_out_raw, max_len_inp,\n",
        "      max_len_targ, vocab_to_int)\n",
        "  input_tensor_val2, target_tensor_val2 = create_dataset_tensors(\n",
        "      test_hard_examples_in_raw, test_hard_examples_out_raw, max_len_inp,\n",
        "      max_len_targ, vocab_to_int)\n",
        "\n",
        "  return (vocab, vocab_to_int,\n",
        "          input_tensor_train, target_tensor_train,\n",
        "          [input_tensor_val0, input_tensor_val1,\n",
        "           input_tensor_val2],\n",
        "          [target_tensor_val0,target_tensor_val1,\n",
        "           target_tensor_val2])\n",
        "\n",
        "\n",
        "def create_cartesian_dataset(trainsize, testsize, vocab, vocab_to_int,\n",
        "                             trainmindigits=1,trainmaxdigits=6,\n",
        "                             testmindigits=7, testmaxdigits=8, reverse=False):\n",
        "  \"\"\"Creates the cartessian product dataset.\"\"\"\n",
        "\n",
        "  def create_example(minlen, maxlen):\n",
        "    symbols1 = [\"0\", \"1\", \"2\", \"3\", \"4\", \"5\", \"6\", \"7\", \"8\", \"9\"]\n",
        "    symbols2 = [\"a\", \"b\", \"c\", \"d\", \"e\", \"f\", \"g\", \"h\", \"i\", \"j\"]\n",
        "    l1 = random.randint(minlen, maxlen)\n",
        "    l2 = random.randint(minlen, maxlen)\n",
        "    set1 = []\n",
        "    set2 = []\n",
        "    for i in range(l1):\n",
        "      set1.append(random.choice(symbols1))\n",
        "    for i in range(l2):\n",
        "      set2.append(random.choice(symbols2))\n",
        "    example_in = set1 + [SEP_TOKEN] + set2 + [END_TOKEN]\n",
        "    example_out = []\n",
        "    for i in set1:\n",
        "      for j in set2:\n",
        "        example_out.append(i)\n",
        "        example_out.append(j)\n",
        "        example_out.append(SEP_TOKEN)\n",
        "    example_out.append(END_TOKEN)\n",
        "    if reverse:\n",
        "      return example_out, [START_TOKEN] + example_in\n",
        "    else:\n",
        "      return example_in, [START_TOKEN] + example_out\n",
        "\n",
        "  def create_examples(n, minlen, maxlen):\n",
        "    examples_in = []\n",
        "    examples_out = []\n",
        "    for i in range(n):\n",
        "      ein, eout = create_example(minlen, maxlen)\n",
        "      examples_in.append(ein)\n",
        "      examples_out.append(eout)\n",
        "    return examples_in, examples_out\n",
        "\n",
        "  vocab = [PAD_TOKEN, SEP_TOKEN, END_TOKEN, START_TOKEN,\n",
        "           \"0\", \"1\", \"2\", \"3\", \"4\", \"5\", \"6\", \"7\", \"8\", \"9\",\n",
        "           \"a\", \"b\", \"c\", \"d\", \"e\", \"f\", \"g\", \"h\", \"i\", \"j\"]\n",
        "  vocab_to_int = {PAD_TOKEN:0, SEP_TOKEN:SEP_TOKEN_IDX,\n",
        "                  END_TOKEN:END_TOKEN_IDX,\n",
        "                  START_TOKEN:START_TOKEN_IDX,\n",
        "                  \"0\":4, \"1\":5, \"2\":6, \"3\":7, \"4\":8,\n",
        "                  \"5\":9, \"6\":10, \"7\":11, \"8\":12, \"9\":13,\n",
        "                  \"a\":14, \"b\":15, \"c\":16, \"d\":17, \"e\":18,\n",
        "                  \"f\":19, \"g\":20, \"h\":21, \"i\":22, \"j\":23}\n",
        "\n",
        "  train_examples_in_raw, train_examples_out_raw = create_examples(\n",
        "      trainsize, trainmindigits, trainmaxdigits)\n",
        "  test_easy_examples_in_raw, test_easy_examples_out_raw = create_examples(\n",
        "      testsize, trainmindigits, trainmaxdigits)\n",
        "  test_hard_examples_in_raw, test_hard_examples_out_raw = create_examples(\n",
        "      testsize, testmindigits, testmaxdigits)\n",
        "\n",
        "  max_len_inp = max(max_length(train_examples_in_raw), max_length(\n",
        "      test_easy_examples_in_raw), max_length(test_hard_examples_in_raw))\n",
        "  max_len_targ = max(max_length(train_examples_out_raw), max_length(\n",
        "      test_easy_examples_out_raw), max_length(test_hard_examples_out_raw))\n",
        "\n",
        "  input_tensor_train, target_tensor_train = create_dataset_tensors(\n",
        "      train_examples_in_raw, train_examples_out_raw, max_len_inp, max_len_targ,\n",
        "      vocab_to_int)\n",
        "  input_tensor_val0, target_tensor_val0 = create_dataset_tensors(\n",
        "      train_examples_in_raw[0:testsize], train_examples_out_raw[0:testsize],\n",
        "      max_len_inp, max_len_targ, vocab_to_int)\n",
        "  input_tensor_val1, target_tensor_val1 = create_dataset_tensors(\n",
        "      test_easy_examples_in_raw, test_easy_examples_out_raw, max_len_inp,\n",
        "      max_len_targ, vocab_to_int)\n",
        "  input_tensor_val2, target_tensor_val2 = create_dataset_tensors(\n",
        "      test_hard_examples_in_raw, test_hard_examples_out_raw, max_len_inp,\n",
        "      max_len_targ, vocab_to_int)\n",
        "\n",
        "  return (vocab, vocab_to_int,\n",
        "          input_tensor_train, target_tensor_train,\n",
        "          [input_tensor_val0, input_tensor_val1,\n",
        "           input_tensor_val2],\n",
        "          [target_tensor_val0,target_tensor_val1,\n",
        "           target_tensor_val2])\n",
        "\n",
        "\n",
        "def create_intersection_dataset(trainsize, testsize, vocab, vocab_to_int,\n",
        "                                trainminelements=1, trainmaxelements=16,\n",
        "                                testminelements=17,testmaxelements=24):\n",
        "  \"\"\"Creates the intersection boolean dataset.\"\"\"\n",
        "\n",
        "  elements = []\n",
        "  for a in [\"a\", \"b\", \"c\", \"d\", \"e\", \"f\", \"g\", \"h\", \"i\", \"j\"]:\n",
        "    for b in [\"0\", \"1\", \"2\", \"3\", \"4\", \"5\", \"6\", \"7\", \"8\", \"9\"]:\n",
        "      elements.append(a+b)\n",
        "\n",
        "  def create_example(minlen, maxlen, label):\n",
        "    l1 = random.randint(minlen, maxlen)\n",
        "    l2 = random.randint(minlen, maxlen)\n",
        "    example_in = []\n",
        "    intersection = []\n",
        "    set1 = []\n",
        "    set2 = []\n",
        "    for i in range(l1):\n",
        "      element = random.choice(elements)\n",
        "      if element not in set1:\n",
        "        set1.append(element)\n",
        "    while len(set2) \u003c l2:\n",
        "      element = random.choice(elements)\n",
        "      if element not in set2:\n",
        "        if element in set1:\n",
        "          if label == \"true\":\n",
        "            intersection.append(element)\n",
        "            set2.append(element)\n",
        "        else:\n",
        "          set2.append(element)\n",
        "    if label == \"true\" and not intersection:\n",
        "      element = random.choice(set1)\n",
        "      set2[random.choice(range(len(set2)))] = element\n",
        "      intersection.append(element)\n",
        "\n",
        "    example_in = set1 + [SEP_TOKEN] + set2 + [END_TOKEN]\n",
        "    if intersection:\n",
        "      example_out = [START_TOKEN, \"true\", END_TOKEN]\n",
        "    else:\n",
        "      example_out = [START_TOKEN, \"false\", END_TOKEN]\n",
        "    return example_in, example_out\n",
        "\n",
        "  def create_examples(n, minlen, maxlen):\n",
        "    examples_in = []\n",
        "    examples_out = []\n",
        "    n_positive = 0\n",
        "    for i in range(n):\n",
        "      ein, eout = create_example(minlen, maxlen, [\"true\", \"false\"][i%2])\n",
        "      examples_in.append(ein)\n",
        "      examples_out.append(eout)\n",
        "      if \"true\" in eout:\n",
        "        n_positive += 1\n",
        "    print(f\"positive: {n_positive}, negative: {n - n_positive}\")\n",
        "    return examples_in, examples_out\n",
        "\n",
        "  vocab = [PAD_TOKEN, SEP_TOKEN, END_TOKEN, START_TOKEN,\n",
        "           \"true\", \"false\"] + elements\n",
        "  vocab_to_int = {PAD_TOKEN:0, SEP_TOKEN:SEP_TOKEN_IDX,\n",
        "                  END_TOKEN:END_TOKEN_IDX,\n",
        "                  START_TOKEN:START_TOKEN_IDX,\n",
        "                  \"true\": 4, \"false\": 5}\n",
        "  for element in elements:\n",
        "    vocab_to_int[element] = len(vocab_to_int)\n",
        "\n",
        "  train_examples_in_raw, train_examples_out_raw = create_examples(\n",
        "      trainsize, trainminelements, trainmaxelements)\n",
        "  test_easy_examples_in_raw, test_easy_examples_out_raw = create_examples(\n",
        "      testsize, trainminelements, trainmaxelements)\n",
        "  test_hard_examples_in_raw, test_hard_examples_out_raw = create_examples(\n",
        "      testsize, testminelements, testmaxelements)\n",
        "\n",
        "  max_len_inp = max(max_length(train_examples_in_raw), max_length(\n",
        "      test_easy_examples_in_raw), max_length(test_hard_examples_in_raw))\n",
        "  max_len_targ = max(max_length(train_examples_out_raw), max_length(\n",
        "      test_easy_examples_out_raw), max_length(test_hard_examples_out_raw))\n",
        "\n",
        "  input_tensor_train, target_tensor_train = create_dataset_tensors(\n",
        "      train_examples_in_raw, train_examples_out_raw, max_len_inp, max_len_targ,\n",
        "      vocab_to_int)\n",
        "  input_tensor_val0, target_tensor_val0 = create_dataset_tensors(\n",
        "      train_examples_in_raw[0:testsize], train_examples_out_raw[0:testsize],\n",
        "      max_len_inp, max_len_targ, vocab_to_int)\n",
        "  input_tensor_val1, target_tensor_val1 = create_dataset_tensors(\n",
        "      test_easy_examples_in_raw, test_easy_examples_out_raw, max_len_inp,\n",
        "      max_len_targ, vocab_to_int)\n",
        "  input_tensor_val2, target_tensor_val2 = create_dataset_tensors(\n",
        "      test_hard_examples_in_raw, test_hard_examples_out_raw, max_len_inp,\n",
        "      max_len_targ, vocab_to_int)\n",
        "\n",
        "  return (vocab, vocab_to_int,\n",
        "          input_tensor_train, target_tensor_train,\n",
        "          [input_tensor_val0, input_tensor_val1,\n",
        "           input_tensor_val2],\n",
        "          [target_tensor_val0,target_tensor_val1,\n",
        "           target_tensor_val2])"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "id": "RsepoMW10OB7"
      },
      "outputs": [],
      "source": [
        "#@title \"SCAN\" dataset generation\n",
        "\n",
        "# Set the appropriate path here:\n",
        "to_download = [\n",
        "               [\"tasks_train_length.txt\",\n",
        "                \"SCAN/length_split/tasks_train_length.txt\"],\n",
        "               [\"tasks_test_length.txt\",\n",
        "                \"SCAN/length_split/tasks_test_length.txt\"],\n",
        "               [\"tasks_train_addprim_jump.txt\",\n",
        "                \"SCAN/add_prim_split/tasks_train_addprim_jump.txt\"],\n",
        "               [\"tasks_test_addprim_jump.txt\",\n",
        "                \"SCAN/add_prim_split/tasks_test_addprim_jump.txt\"],\n",
        "               ]\n",
        "uploaded = {}\n",
        "\n",
        "# Preload the files:\n",
        "for [name, path] in to_download:\n",
        "  with tf.io.gfile.GFile(path, \"rb\") as f:\n",
        "    uploaded[name] = f.read()\n",
        "    lines = uploaded[name].decode(\"utf-8\").split(\"\\n\")\n",
        "\n",
        "def create_scan_length_dataset(vocab, vocab_to_int,\n",
        "                               map_output_to_input=False):\n",
        "  return create_scan_dataset(vocab, vocab_to_int,\n",
        "                              \"tasks_train_length.txt\",\n",
        "                              \"tasks_test_length.txt\",\n",
        "                              map_output_to_input=map_output_to_input)\n",
        "\n",
        "def create_scan_add_jump_dataset(vocab, vocab_to_int,\n",
        "                             map_output_to_input=False):\n",
        "  return create_scan_dataset(vocab, vocab_to_int,\n",
        "                              \"tasks_train_addprim_jump.txt\",\n",
        "                              \"tasks_test_addprim_jump.txt\",\n",
        "                              map_output_to_input=map_output_to_input)\n",
        "\n",
        "def create_scan_dataset(vocab, vocab_to_int,\n",
        "                         train_filename,\n",
        "                         test_filename,\n",
        "                         map_output_to_input=False,\n",
        "                         keep_all_steps=False):\n",
        "  \"\"\"Creates a version of the SCAN dataset.\"\"\"\n",
        "\n",
        "  vocab = [PAD_TOKEN, SEP_TOKEN, END_TOKEN, START_TOKEN,\n",
        "          END_ITERATION_TOKEN]\n",
        "  vocab_to_int = {PAD_TOKEN:0, SEP_TOKEN:SEP_TOKEN_IDX,\n",
        "                  END_TOKEN:END_TOKEN_IDX,\n",
        "                  START_TOKEN:START_TOKEN_IDX,\n",
        "                  END_ITERATION_TOKEN:END_ITERATION_TOKEN_IDX}\n",
        "\n",
        "  def create_dataset_tensors_scan(instances, maxlen_inp=None, maxlen_targ=None,\n",
        "                                  keep_all_steps=False):\n",
        "    in_tensor = []\n",
        "    out_tensor = []\n",
        "    for instance in instances:\n",
        "      if keep_all_steps:\n",
        "        for i in range(len(instance)-1):\n",
        "          in_tensor.append(instance[i])\n",
        "          out_tensor.append(instance[i+1])\n",
        "      else:\n",
        "        # keep only the first and last steps (ignore intermediate steps):\n",
        "        in_tensor.append(instance[0])\n",
        "        out_tensor.append(instance[-1])\n",
        "\n",
        "    in_tensor = tf.keras.preprocessing.sequence.pad_sequences(\n",
        "        in_tensor,padding='post', maxlen=maxlen_inp)\n",
        "    out_tensor = tf.keras.preprocessing.sequence.pad_sequences(\n",
        "        out_tensor, padding='post', maxlen=maxlen_targ)\n",
        "\n",
        "    return in_tensor, out_tensor\n",
        "\n",
        "  def tokenize_scan_line(line, vocab, vocab_to_int, map_output_to_input):\n",
        "    instance_in_split = line.split(\" \")\n",
        "\n",
        "    if map_output_to_input:\n",
        "      for i in range(len(instance_in_split)):\n",
        "        if instance_in_split[i] == \"I_WALK\":\n",
        "          instance_in_split[i] = \"walk\"\n",
        "        if instance_in_split[i] == \"I_JUMP\":\n",
        "          instance_in_split[i] = \"jump\"\n",
        "        if instance_in_split[i] == \"I_LOOK\":\n",
        "          instance_in_split[i] = \"look\"\n",
        "        if instance_in_split[i] == \"I_RUN\":\n",
        "          instance_in_split[i] = \"run\"\n",
        "\n",
        "    # find the tokens:\n",
        "    for token in instance_in_split:\n",
        "      if token not in vocab_to_int:\n",
        "        vocab_to_int[token] = len(vocab)\n",
        "        vocab.append(token)\n",
        "\n",
        "    # tokenize:\n",
        "    instance_in_tokenized = (\n",
        "        [START_TOKEN_IDX] +\n",
        "        [vocab_to_int[x] for x in instance_in_split] +\n",
        "        [END_TOKEN_IDX])\n",
        "    return instance_in_tokenized\n",
        "\n",
        "  def load_and_tokenize_data(filename):\n",
        "    instances_raw = []\n",
        "    instances = []\n",
        "    lines = uploaded[filename].decode(\"utf-8\").split(\"\\n\")\n",
        "    for line in lines:\n",
        "      if line.startswith(\"IN:\"):\n",
        "        line = line[4:]\n",
        "        instance_raw = line.split(\" OUT: \")\n",
        "        instance = [tokenize_scan_line(instance_raw[0], vocab, vocab_to_int,\n",
        "                                       map_output_to_input),\n",
        "                    tokenize_scan_line(instance_raw[1], vocab, vocab_to_int,\n",
        "                                       map_output_to_input)]\n",
        "        instances_raw.append(instance_raw)\n",
        "        instances.append(instance)\n",
        "\n",
        "    print(\"# instances: \" + str(len(instances)))\n",
        "    return instances_raw, instances\n",
        "\n",
        "  instances_train_raw, instances_train = load_and_tokenize_data(train_filename)\n",
        "  instances_test_raw, instances_test = load_and_tokenize_data(test_filename)\n",
        "\n",
        "  input_tensor_train, target_tensor_train = create_dataset_tensors_scan(\n",
        "      instances_train, keep_all_steps=keep_all_steps)\n",
        "  input_tensor_val, target_tensor_val = create_dataset_tensors_scan(\n",
        "      instances_test, keep_all_steps=keep_all_steps)\n",
        "  max_length_train = max_length(input_tensor_train)\n",
        "  max_length_val = max_length(input_tensor_val)\n",
        "  max_length_targ_train = max_length(target_tensor_train)\n",
        "  max_length_targ_val = max_length(target_tensor_val)\n",
        "\n",
        "  testsize = len(instances_test)\n",
        "  max_len_inp = max(max_length_train, max_length_val)\n",
        "  max_len_targ = max(max_length_targ_train, max_length_targ_val)\n",
        "  input_tensor_train, target_tensor_train = create_dataset_tensors_scan(\n",
        "      instances_train, maxlen_inp=max_len_inp, maxlen_targ=max_len_targ,\n",
        "      keep_all_steps=keep_all_steps)\n",
        "  input_tensor_val0, target_tensor_val0 = create_dataset_tensors_scan(\n",
        "      instances_train[0:testsize], maxlen_inp=max_len_inp,\n",
        "      maxlen_targ=max_len_targ, keep_all_steps=keep_all_steps)\n",
        "  input_tensor_val, target_tensor_val = create_dataset_tensors_scan(\n",
        "      instances_test, maxlen_inp=max_len_inp, maxlen_targ=max_len_targ,\n",
        "      keep_all_steps=keep_all_steps)\n",
        "\n",
        "  return (vocab, vocab_to_int,\n",
        "          input_tensor_train, target_tensor_train,\n",
        "          [input_tensor_val0, input_tensor_val],\n",
        "          [target_tensor_val0, target_tensor_val])\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "id": "9YtpcTOQq9Kh"
      },
      "outputs": [],
      "source": [
        "#@title \"PCFG\" dataset generation\n",
        "\n",
        "# Set the appropriate path here:\n",
        "to_download_pcfg = [\n",
        "                    [\"pcfg_productivity_train\",\n",
        "                     \"pcfg/productivity/train.src\",\n",
        "                     \"pcfg/productivity/train.tgt\"],\n",
        "                    [\"pcfg_productivity_test\",\n",
        "                     \"pcfg/productivity/test.src\",\n",
        "                     \"pcfg/productivity/test.tgt\"],\n",
        "\n",
        "                    [\"pcfg_systematicity_train\",\n",
        "                     \"pcfg/systematicity/train.src\",\n",
        "                     \"pcfg/systematicity/train.tgt\"],\n",
        "                    [\"pcfg_systematicity_test\",\n",
        "                     \"pcfg/systematicity/test.src\",\n",
        "                     \"pcfg/systematicity/test.tgt\"],\n",
        "                   ]\n",
        "uploaded_pcfg = {}\n",
        "\n",
        "MAX_TRAIN_LEN = 128\n",
        "MAX_TEST_LEN = 256\n",
        "\n",
        "# Preloading the PCFG files:\n",
        "for [name, pathin, pathout] in to_download_pcfg:\n",
        "  with (tf.io.gfile.GFile(pathin, \"rb\") as fin,\n",
        "        tf.io.gfile.GFile(pathout, \"rb\") as fout):\n",
        "    uploaded_pcfg[name] = [fin.read(), fout.read()]\n",
        "    lines = uploaded_pcfg[name][0].decode(\"utf-8\").split(\"\\n\")\n",
        "    print(name + \": \" + str(len(lines)) + \" lines\")\n",
        "\n",
        "\n",
        "def create_pcfg_dataset(pcfg_split, vocab, vocab_to_int):\n",
        "  \"\"\"Creates a version of the PCFG dataset.\"\"\"\n",
        "\n",
        "  vocab = [PAD_TOKEN, SEP_TOKEN, END_TOKEN, START_TOKEN,\n",
        "          END_ITERATION_TOKEN]\n",
        "  vocab_to_int = {PAD_TOKEN:0, SEP_TOKEN:SEP_TOKEN_IDX,\n",
        "                  END_TOKEN:END_TOKEN_IDX,\n",
        "                  START_TOKEN:START_TOKEN_IDX,\n",
        "                  END_ITERATION_TOKEN:END_ITERATION_TOKEN_IDX}\n",
        "\n",
        "  def create_dataset_tensors_pcfg(instances, maxlen_inp=None, maxlen_targ=None):\n",
        "    in_tensor = []\n",
        "    out_tensor = []\n",
        "    for instance in instances:\n",
        "      for i in range(len(instance)-1):\n",
        "        in_tensor.append(instance[i])\n",
        "        # out_tensor.append(instance[i+1] + [END_ITERATION_TOKEN])\n",
        "        out_tensor.append(instance[i+1])\n",
        "\n",
        "    in_tensor = tf.keras.preprocessing.sequence.pad_sequences(\n",
        "        in_tensor, padding='post', maxlen=maxlen_inp)\n",
        "    out_tensor = tf.keras.preprocessing.sequence.pad_sequences(\n",
        "        out_tensor, padding='post', maxlen=maxlen_targ)\n",
        "\n",
        "    return in_tensor, out_tensor\n",
        "\n",
        "  def load_and_tokenize_data(filename, maxlen):\n",
        "    max_in_len = 0\n",
        "    max_out_len = 0\n",
        "    instances_raw = []\n",
        "    instances = []\n",
        "    lines_in = uploaded_pcfg[filename][0].decode(\"utf-8\").split(\"\\n\")\n",
        "    lines_out = uploaded_pcfg[filename][1].decode(\"utf-8\").split(\"\\n\")\n",
        "    instance_raw = []\n",
        "    instance = []\n",
        "    for i in range(len(lines_in)):\n",
        "      instance_raw = [lines_in[i], lines_out[i]]\n",
        "      instances_raw.append(instance_raw)\n",
        "\n",
        "      for instance_part in instance_raw:\n",
        "        for token in instance_part.split(\" \"):\n",
        "          if token not in vocab_to_int:\n",
        "            vocab_to_int[token] = len(vocab)\n",
        "            vocab.append(token)\n",
        "\n",
        "      # tokenize:\n",
        "      instance_in_tokenized = (\n",
        "          [START_TOKEN_IDX] +\n",
        "          [vocab_to_int[x] for x in instance_raw[0].split(\" \")] +\n",
        "          [END_TOKEN_IDX])\n",
        "      instance_out_tokenized = (\n",
        "          [START_TOKEN_IDX] +\n",
        "          [vocab_to_int[x] for x in instance_raw[1].split(\" \")] +\n",
        "          [END_TOKEN_IDX])\n",
        "      if len(instance_out_tokenized) \u003e maxlen:\n",
        "        continue\n",
        "      instances.append([instance_in_tokenized, instance_out_tokenized])\n",
        "\n",
        "      max_in_len = max(max_in_len, len(instance_in_tokenized))\n",
        "      max_out_len = max(max_out_len, len(instance_out_tokenized))\n",
        "\n",
        "    print(\"max_in_len: \" + str(max_in_len))\n",
        "    print(\"max_out_len: \" + str(max_out_len))\n",
        "    return instances_raw, instances\n",
        "\n",
        "  instances_train_raw, instances_train = load_and_tokenize_data(\n",
        "      \"pcfg_\"+pcfg_split+\"_train\", MAX_TRAIN_LEN)\n",
        "  instances_test_raw, instances_test = load_and_tokenize_data(\n",
        "      \"pcfg_\"+pcfg_split+\"_test\", MAX_TEST_LEN)\n",
        "\n",
        "  input_tensor_train, target_tensor_train = create_dataset_tensors_pcfg(\n",
        "      instances_train)\n",
        "  input_tensor_val, target_tensor_val = create_dataset_tensors_pcfg(\n",
        "      instances_test)\n",
        "  max_length_train = max_length(input_tensor_train)\n",
        "  max_length_val = max_length(input_tensor_val)\n",
        "  max_length_targ_train = max_length(target_tensor_train)\n",
        "  max_length_targ_val = max_length(target_tensor_val)\n",
        "\n",
        "  testset_size = len(instances_test)\n",
        "  max_len_inp = max(max_length_train, max_length_val)\n",
        "  max_len_targ = max(max_length_targ_train, max_length_targ_val)\n",
        "  input_tensor_train, target_tensor_train = create_dataset_tensors_pcfg(\n",
        "      instances_train, maxlen_inp=max_len_inp, maxlen_targ=max_len_targ)\n",
        "  input_tensor_val0, target_tensor_val0 = create_dataset_tensors_pcfg(\n",
        "      instances_train[0:testset_size], maxlen_inp=max_len_inp,\n",
        "      maxlen_targ=max_len_targ)\n",
        "  input_tensor_val, target_tensor_val = create_dataset_tensors_pcfg(\n",
        "      instances_test, maxlen_inp=max_len_inp, maxlen_targ=max_len_targ)\n",
        "\n",
        "  return (vocab, vocab_to_int,\n",
        "          input_tensor_train, target_tensor_train,\n",
        "          [input_tensor_val0, input_tensor_val],\n",
        "          [target_tensor_val0, target_tensor_val])"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "id": "ci9LM2RxLCys"
      },
      "outputs": [],
      "source": [
        "#@title \"COGS\" dataset generation\n",
        "\n",
        "# Set the appropriate path here:\n",
        "to_download_cogs = [\n",
        "                    [\"COGS_train\",\n",
        "                     \"COGS/data/train.tsv\"],\n",
        "                    [\"COGS_test\",\n",
        "                     \"COGS/data/test.tsv\"],\n",
        "                    [\"COGS_gen\",\n",
        "                     \"COGS/data/gen.tsv\"],\n",
        "                   ]\n",
        "uploaded_cogs = {}\n",
        "\n",
        "# Preloading all the COGS files:\n",
        "for [name, path] in to_download_cogs:\n",
        "  with tf.io.gfile.GFile(path, \"rb\") as f:\n",
        "    uploaded_cogs[name] = f.read()\n",
        "    lines = uploaded_cogs[name].decode(\"utf-8\").split(\"\\n\")\n",
        "    print(name + \": \" + str(len(lines)) + \" lines\")\n",
        "\n",
        "def create_cogs_dataset(vocab, vocab_to_int, max_len=256):\n",
        "  \"\"\"Generates the COGS dataset.\n",
        "\n",
        "  Note: set 'max_len' to 512 to ensure even the longest instances fit.\n",
        "  Otherwise, longer instances will be discarded. The function that loads this\n",
        "  dataset via the 'DATASET_COGS_FULL' name (used in the paper) already does\n",
        "  this.\"\"\"\n",
        "\n",
        "  MAX_COGS_TRAIN_LEN = max_len\n",
        "  MAX_COGS_TEST_LEN = max_len\n",
        "\n",
        "  vocab = [PAD_TOKEN, SEP_TOKEN, END_TOKEN, START_TOKEN,\n",
        "          END_ITERATION_TOKEN]\n",
        "  vocab_to_int = {PAD_TOKEN:0, SEP_TOKEN:SEP_TOKEN_IDX,\n",
        "                  END_TOKEN:END_TOKEN_IDX,\n",
        "                  START_TOKEN:START_TOKEN_IDX,\n",
        "                  END_ITERATION_TOKEN:END_ITERATION_TOKEN_IDX}\n",
        "\n",
        "  def create_dataset_tensors_cogs(instances, maxlen_inp=None, maxlen_targ=None):\n",
        "    in_tensor = []\n",
        "    out_tensor = []\n",
        "    for instance in instances:\n",
        "      for i in range(len(instance)-1):\n",
        "        in_tensor.append(instance[i])\n",
        "        # out_tensor.append(instance[i+1] + [END_ITERATION_TOKEN])\n",
        "        out_tensor.append(instance[i+1])\n",
        "\n",
        "    in_tensor = tf.keras.preprocessing.sequence.pad_sequences(\n",
        "        in_tensor, padding='post', maxlen=maxlen_inp)\n",
        "    out_tensor = tf.keras.preprocessing.sequence.pad_sequences(\n",
        "        out_tensor, padding='post', maxlen=maxlen_targ)\n",
        "\n",
        "    return in_tensor, out_tensor\n",
        "\n",
        "  def load_and_tokenize_data_by_distribution(filename, maxlen):\n",
        "    max_in_len = 0\n",
        "    max_out_len = 0\n",
        "    instances_raw = []\n",
        "    instances = []\n",
        "    lines = uploaded_cogs[filename].decode(\"utf-8\").split(\"\\n\")\n",
        "    distribution_names = []\n",
        "    lines_by_distribution = {}\n",
        "    for line in lines:\n",
        "      parts = line.split(\"\\t\")\n",
        "      if len(parts) == 3:\n",
        "        distribution = parts[2]\n",
        "        if distribution in lines_by_distribution:\n",
        "          lines_by_distribution[distribution].append(line)\n",
        "        else:\n",
        "          distribution_names.append(distribution)\n",
        "          lines_by_distribution[distribution] = [line]\n",
        "\n",
        "    for distribution in distribution_names:\n",
        "      print(f\"{distribution}: {len(lines_by_distribution[distribution])}\")\n",
        "\n",
        "    instances_raw_by_distribution = []\n",
        "    instances_by_distribution = []\n",
        "    for distribution in distribution_names:\n",
        "      intances_raw, instances = load_and_tokenize_data_internal(\n",
        "          lines_by_distribution[distribution], maxlen)\n",
        "      instances_raw_by_distribution.append(intances_raw)\n",
        "      instances_by_distribution.append(instances)\n",
        "    return instances_raw_by_distribution, instances_by_distribution\n",
        "\n",
        "\n",
        "  def load_and_tokenize_data(filename, maxlen):\n",
        "    lines = uploaded_cogs[filename].decode(\"utf-8\").split(\"\\n\")\n",
        "    return load_and_tokenize_data_internal(lines, maxlen)\n",
        "\n",
        "\n",
        "  def load_and_tokenize_data_internal(lines, maxlen):\n",
        "    max_in_len = 0\n",
        "    max_out_len = 0\n",
        "    instances_raw = []\n",
        "    instances = []\n",
        "    lines_in = []\n",
        "    lines_out = []\n",
        "    for line in lines:\n",
        "      parts = line.split(\"\\t\")\n",
        "      if len(parts) == 3:\n",
        "        lines_in.append(parts[0])\n",
        "        lines_out.append(parts[1])\n",
        "    instance_raw = []\n",
        "    instance = []\n",
        "    for i in range(len(lines_in)):\n",
        "      instance_raw = [lines_in[i], lines_out[i]]\n",
        "      instances_raw.append(instance_raw)\n",
        "\n",
        "      for instance_part in instance_raw:\n",
        "        for token in instance_part.split(\" \"):\n",
        "          if token not in vocab_to_int:\n",
        "            vocab_to_int[token] = len(vocab)\n",
        "            vocab.append(token)\n",
        "\n",
        "      # tokenize:\n",
        "      instance_in_tokenized = (\n",
        "          [START_TOKEN_IDX] +\n",
        "          [vocab_to_int[x] for x in instance_raw[0].split(\" \")] +\n",
        "          [END_TOKEN_IDX])\n",
        "      instance_out_tokenized = (\n",
        "          [START_TOKEN_IDX] +\n",
        "          [vocab_to_int[x] for x in instance_raw[1].split(\" \")] +\n",
        "          [END_TOKEN_IDX])\n",
        "      if len(instance_out_tokenized) \u003e maxlen:\n",
        "        continue\n",
        "      instances.append([instance_in_tokenized, instance_out_tokenized])\n",
        "\n",
        "      max_in_len = max(max_in_len, len(instance_in_tokenized))\n",
        "      max_out_len = max(max_out_len, len(instance_out_tokenized))\n",
        "\n",
        "    print(\"max_in_len: \" + str(max_in_len))\n",
        "    print(\"max_out_len: \" + str(max_out_len))\n",
        "    return instances_raw, instances\n",
        "\n",
        "  instances_train_raw, instances_train = load_and_tokenize_data(\n",
        "      \"COGS_train\", MAX_COGS_TRAIN_LEN)\n",
        "  instances_test_raw, instances_test = load_and_tokenize_data(\n",
        "      \"COGS_test\", MAX_COGS_TEST_LEN)\n",
        "  instances_gen_raw, instances_gen = load_and_tokenize_data(\n",
        "      \"COGS_gen\", MAX_COGS_TEST_LEN)\n",
        "  instances_gen_raw_by_dist, instances_gen_by_dist = (\n",
        "      load_and_tokenize_data_by_distribution(\"COGS_gen\", MAX_COGS_TEST_LEN))\n",
        "\n",
        "  instances_list = [instances_train, instances_train[:len(instances_test)],\n",
        "                    instances_test] + instances_gen_by_dist + [instances_gen]\n",
        "\n",
        "  input_tensors = []\n",
        "  target_tensors = []\n",
        "  max_len_inp = 0\n",
        "  max_len_targ = 0\n",
        "  for instances in instances_list:\n",
        "    input_tensor, target_tensor = create_dataset_tensors_cogs(instances)\n",
        "    input_tensors.append(input_tensor)\n",
        "    target_tensors.append(target_tensor)\n",
        "    max_len_inp = max(max_len_inp, max_length(input_tensor))\n",
        "    max_len_targ = max(max_len_targ, max_length(input_tensor))\n",
        "\n",
        "  input_tensors = []\n",
        "  target_tensors = []\n",
        "  for instances in instances_list:\n",
        "    input_tensor, target_tensor = create_dataset_tensors_cogs(\n",
        "        instances, maxlen_inp=max_len_inp, maxlen_targ=max_len_targ)\n",
        "    input_tensors.append(input_tensor)\n",
        "    target_tensors.append(target_tensor)\n",
        "    max_len_inp = max(max_len_inp, max_length(input_tensor))\n",
        "    max_len_targ = max(max_len_targ, max_length(input_tensor))\n",
        "\n",
        "  return (vocab, vocab_to_int,\n",
        "          input_tensors[0], target_tensors[0],\n",
        "          input_tensors[1:], target_tensors[1:])\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "id": "OV68o8dRCAIy"
      },
      "outputs": [],
      "source": [
        "#@title \"CFQ\" dataset generation\n",
        "\n",
        "# Set the appropriate path here:\n",
        "to_download_cfq = [\n",
        "                    [\"cfq_mcd1_train\",\n",
        "                     \"cfq/mcd1/train.txt\"],\n",
        "                    [\"cfq_mcd1_dev\",\n",
        "                     \"cfq/mcd1/dev.txt\"],\n",
        "                    [\"cfq_mcd1_test\",\n",
        "                     \"cfq/mcd1/test.txt\"],\n",
        "                   ]\n",
        "uploaded_cfq = {}\n",
        "\n",
        "MAX_TRAIN_LEN = 256\n",
        "MAX_TEST_LEN = 256\n",
        "\n",
        "# simulate uploading them from local hard drive:\n",
        "for [name, path] in to_download_cfq:\n",
        "  with tf.io.gfile.GFile(path, \"rb\") as f:\n",
        "    uploaded_cfq[name] = f.read()\n",
        "    lines = uploaded_cfq[name].decode(\"utf-8\").split(\"\\n\")\n",
        "    print(name + \": \" + str(len(lines)) + \" lines\")\n",
        "\n",
        "\n",
        "def cfq_decompose_output(line):\n",
        "  tokens = line.split(\" \")\n",
        "  prefix = \"\"\n",
        "  postfix = \"\"\n",
        "  triplets_text = \"\"\n",
        "  state = 0\n",
        "  for token in tokens:\n",
        "    if state == 0:\n",
        "      if token == \"{\":\n",
        "        prefix += token + \" \"\n",
        "        state = 1\n",
        "      else:\n",
        "        prefix += token + \" \"\n",
        "    elif state == 1:\n",
        "      if token == \"}\":\n",
        "        postfix += token + \" \"\n",
        "        state = 2\n",
        "      else:\n",
        "        triplets_text += token + \" \"\n",
        "    else:\n",
        "      postfix += token + \" \"\n",
        "  triplets = triplets_text.strip().split(\" . \")\n",
        "  return prefix, triplets, postfix\n",
        "\n",
        "\n",
        "def cfq_rewrite_cartesian(triplets):\n",
        "  if not triplets:\n",
        "    return triplets\n",
        "  triplet = triplets[0]\n",
        "  tokens = triplet.split(\" \")\n",
        "  if len(tokens) == 3 and tokens[1] != \"a\":\n",
        "    relation = tokens[1]\n",
        "    left_tokens = [tokens[0]]\n",
        "    right_tokens = [tokens[2]]\n",
        "    relation_pairs = [(tokens[0], tokens[2])]\n",
        "    to_delete = [triplet]\n",
        "    to_keep = []\n",
        "    for triplet2 in triplets[1:]:\n",
        "      tokens2 = triplet2.split(\" \")\n",
        "      if len(tokens2) == 3 and tokens2[1] == relation:\n",
        "        relation_pairs.append((tokens2[0], tokens2[2]))\n",
        "        if tokens2[0] not in left_tokens:\n",
        "          left_tokens.append(tokens2[0])\n",
        "        if tokens2[2] not in right_tokens:\n",
        "          right_tokens.append(tokens2[2])\n",
        "        to_delete.append(triplet2)\n",
        "      else:\n",
        "        to_keep.append(triplet2)\n",
        "    # See if it's a cartesian product:\n",
        "    any_missing = False\n",
        "    for left_token in left_tokens:\n",
        "      for right_token in right_tokens:\n",
        "        if (left_token, right_token) not in relation_pairs:\n",
        "          any_missing = True\n",
        "          break\n",
        "      if any_missing:\n",
        "        break\n",
        "    if any_missing:\n",
        "      return [\"( \" + tokens[0] + \" ) ( \" + relation + \" ) ( \" + tokens[2] +\n",
        "              \" )\"] + cfq_rewrite_cartesian(triplets[1:])\n",
        "    else:\n",
        "      # we have a cartesian product!\n",
        "      new_triplet = (\"( \" + \" \".join(left_tokens) + \" ) ( \" + relation + \" ) ( \"\n",
        "      + \" \".join(right_tokens) + \" )\")\n",
        "      return [new_triplet] + cfq_rewrite_cartesian(to_keep)\n",
        "\n",
        "  else:\n",
        "    return [triplet] + cfq_rewrite_cartesian(triplets[1:])\n",
        "\n",
        "\n",
        "def cfq_merge_cartesians(triplets):\n",
        "  if not triplets:\n",
        "    return triplets\n",
        "  triplet = triplets[0]\n",
        "  if triplet[0] == \"(\":\n",
        "    tokens = triplet.split(\" ) ( \")\n",
        "    if len(tokens) == 3:\n",
        "      to_keep = []\n",
        "      relations = [tokens[1]]\n",
        "      for triplet2 in triplets[1:]:\n",
        "        if triplet2[0] == \"(\":\n",
        "          tokens2 = triplet2.split(\" ) ( \")\n",
        "          if (len(tokens2) == 3 and tokens[0] == tokens2[0] and\n",
        "              tokens[2] == tokens2[2]):\n",
        "            relations.append(tokens2[1])\n",
        "          else:\n",
        "            to_keep.append(triplet2)\n",
        "        else:\n",
        "          to_keep.append(triplet2)\n",
        "      new_triplet = (tokens[0] + \" ) ( \" + \" \".join(relations) + \" ) ( \" +\n",
        "                     tokens[2]);\n",
        "      return [new_triplet] + cfq_merge_cartesians(to_keep)\n",
        "    else:\n",
        "      return [triplet] + cfq_merge_cartesians(triplets[1:])\n",
        "  else:\n",
        "    return [triplet] + cfq_merge_cartesians(triplets[1:])\n",
        "\n",
        "\n",
        "def simplify_cfq_output(output):\n",
        "  prefix, triplets, postfix = cfq_decompose_output(output)\n",
        "  triplets = cfq_rewrite_cartesian(triplets)\n",
        "  triplets = cfq_merge_cartesians(triplets)\n",
        "  return prefix + \" . \".join(triplets) + \" \" + postfix\n",
        "\n",
        "\n",
        "def expand_cfq_output(output):\n",
        "  prefix, triplets, postfix = cfq_decompose_output(output)\n",
        "  new_triplets = []\n",
        "  for triplet in triplets:\n",
        "    if triplet[0] == \"(\":\n",
        "      tokens = triplet.split(\" ) ( \")\n",
        "      left_elements = tokens[0][2:].split(\" \")\n",
        "      relations = tokens[1].split(\" \")\n",
        "      right_elements = tokens[2][:-2].split(\" \")\n",
        "      for left_element in left_elements:\n",
        "        for relation in relations:\n",
        "          for right_element in right_elements:\n",
        "            new_triplets.append(left_element + \" \" + relation + \" \" +\n",
        "                                right_element)\n",
        "    else:\n",
        "      new_triplets.append(triplet)\n",
        "  return prefix + \" . \".join(new_triplets) + \" \" + postfix\n",
        "\n",
        "\n",
        "def cfq_equivalent_output(output1, output2):\n",
        "  prefix1, triplets1, postfix1 = cfq_decompose_output(output1.strip())\n",
        "  prefix2, triplets2, postfix2 = cfq_decompose_output(output2.strip())\n",
        "  if prefix1 != prefix2:\n",
        "    return False\n",
        "  if postfix1 != postfix2:\n",
        "    return False\n",
        "  if len(triplets1) != len(triplets2):\n",
        "    return False\n",
        "  for triplet in triplets1:\n",
        "    if triplet not in triplets2:\n",
        "      return False\n",
        "  for triplet in triplets2:\n",
        "    if triplet not in triplets1:\n",
        "      return False\n",
        "  return True\n",
        "\n",
        "\n",
        "def create_cfq_mcd1_dataset(vocab, vocab_to_int, simplify_cartesians=False):\n",
        "  \"\"\"Creates a version of the CFQ dataset.\"\"\"\n",
        "\n",
        "  vocab = [PAD_TOKEN, SEP_TOKEN, END_TOKEN, START_TOKEN,\n",
        "          END_ITERATION_TOKEN]\n",
        "  vocab_to_int = {PAD_TOKEN:0, SEP_TOKEN:SEP_TOKEN_IDX,\n",
        "                  END_TOKEN:END_TOKEN_IDX,\n",
        "                  START_TOKEN:START_TOKEN_IDX,\n",
        "                  END_ITERATION_TOKEN:END_ITERATION_TOKEN_IDX}\n",
        "\n",
        "  def create_dataset_tensors_cfq(instances, maxlen_inp=None, maxlen_targ=None):\n",
        "    in_tensor = []\n",
        "    out_tensor = []\n",
        "    for instance in instances:\n",
        "      for i in range(len(instance)-1):\n",
        "        in_tensor.append(instance[i])\n",
        "        out_tensor.append(instance[i+1])\n",
        "\n",
        "    in_tensor = tf.keras.preprocessing.sequence.pad_sequences(\n",
        "        in_tensor, padding='post', maxlen=maxlen_inp)\n",
        "    out_tensor = tf.keras.preprocessing.sequence.pad_sequences(\n",
        "        out_tensor, padding='post', maxlen=maxlen_targ)\n",
        "\n",
        "    return in_tensor, out_tensor\n",
        "\n",
        "  def load_and_tokenize_data(filename, maxlen):\n",
        "    max_in_len = 0\n",
        "    max_out_len = 0\n",
        "    instances_raw = []\n",
        "    instances = []\n",
        "    lines = uploaded_cfq[filename].decode(\"utf-8\").split(\"\\n\")\n",
        "    lines_in = []\n",
        "    lines_out = []\n",
        "    for i in range(len(lines)//2):\n",
        "      lines_in.append(lines[i*2].strip().replace(\"INPUT: \", \"\"))\n",
        "      lines_out.append(lines[i*2+1].strip().replace(\"OUTPUT: \", \"\"))\n",
        "    instance_raw = []\n",
        "    instance = []\n",
        "    for i in range(len(lines_in)):\n",
        "      instance_raw = [lines_in[i], lines_out[i]]\n",
        "      if simplify_cartesians:\n",
        "        instance_raw[1] = simplify_cfq_output(lines_out[i])\n",
        "      instances_raw.append(instance_raw)\n",
        "\n",
        "      for instance_part in instance_raw:\n",
        "        for token in instance_part.split(\" \"):\n",
        "          if token not in vocab_to_int:\n",
        "            vocab_to_int[token] = len(vocab)\n",
        "            vocab.append(token)\n",
        "\n",
        "      # tokenize:\n",
        "      instance_in_tokenized = (\n",
        "          [START_TOKEN_IDX] +\n",
        "          [vocab_to_int[x] for x in instance_raw[0].split(\" \")] +\n",
        "          [END_TOKEN_IDX])\n",
        "      instance_out_tokenized = (\n",
        "          [START_TOKEN_IDX] +\n",
        "          [vocab_to_int[x] for x in instance_raw[1].split(\" \")] +\n",
        "          [END_TOKEN_IDX])\n",
        "      if len(instance_out_tokenized) \u003e maxlen:\n",
        "        continue\n",
        "      instances.append([instance_in_tokenized, instance_out_tokenized])\n",
        "\n",
        "      max_in_len = max(max_in_len, len(instance_in_tokenized))\n",
        "      max_out_len = max(max_out_len, len(instance_out_tokenized))\n",
        "\n",
        "    print(\"max_in_len: \" + str(max_in_len))\n",
        "    print(\"max_out_len: \" + str(max_out_len))\n",
        "    return instances_raw, instances\n",
        "\n",
        "  instances_train_raw, instances_train = load_and_tokenize_data(\n",
        "      \"cfq_mcd1_train\", MAX_TRAIN_LEN)\n",
        "  instances_test_raw, instances_test = load_and_tokenize_data(\n",
        "      \"cfq_mcd1_dev\", MAX_TEST_LEN)\n",
        "  instances_gen_raw, instances_gen = load_and_tokenize_data(\n",
        "      \"cfq_mcd1_test\", MAX_TEST_LEN)\n",
        "\n",
        "  input_tensor_train, target_tensor_train = create_dataset_tensors_cfq(\n",
        "      instances_train)\n",
        "  input_tensor_val, target_tensor_val = create_dataset_tensors_cfq(\n",
        "      instances_test)\n",
        "  input_tensor_val2, target_tensor_val2 = create_dataset_tensors_cfq(\n",
        "      instances_gen)\n",
        "  max_length_train = max_length(input_tensor_train)\n",
        "  max_length_val = max_length(input_tensor_val)\n",
        "  max_length_val2 = max_length(input_tensor_val2)\n",
        "  max_length_targ_train = max_length(target_tensor_train)\n",
        "  max_length_targ_val = max_length(target_tensor_val)\n",
        "  max_length_targ_val2 = max_length(target_tensor_val2)\n",
        "\n",
        "  testset_size = len(instances_test)\n",
        "  max_len_inp = max(max_length_train, max_length_val, max_length_val2)\n",
        "  max_len_targ = max(max_length_targ_train, max_length_targ_val,\n",
        "                     max_length_targ_val2)\n",
        "  input_tensor_train, target_tensor_train = create_dataset_tensors_cfq(\n",
        "      instances_train, maxlen_inp=max_len_inp, maxlen_targ=max_len_targ)\n",
        "  input_tensor_val0, target_tensor_val0 = create_dataset_tensors_cfq(\n",
        "      instances_train[0:testset_size], maxlen_inp=max_len_inp,\n",
        "      maxlen_targ=max_len_targ)\n",
        "  input_tensor_val, target_tensor_val = create_dataset_tensors_cfq(\n",
        "      instances_test, maxlen_inp=max_len_inp, maxlen_targ=max_len_targ)\n",
        "  input_tensor_val2, target_tensor_val2 = create_dataset_tensors_cfq(\n",
        "      instances_gen, maxlen_inp=max_len_inp, maxlen_targ=max_len_targ)\n",
        "\n",
        "  return (vocab, vocab_to_int,\n",
        "          input_tensor_train, target_tensor_train,\n",
        "          [input_tensor_val0, input_tensor_val, input_tensor_val2],\n",
        "          [target_tensor_val0, target_tensor_val, target_tensor_val2])"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "id": "gjcFDxENhpzQ"
      },
      "outputs": [],
      "source": [
        "#@title This code block defines Transformer Encoder/Decoder layers, following this tutorial (adding relative attention, a copy decoder, and a few other improvements): https://www.tensorflow.org/tutorials/text/transformer\n",
        "\n",
        "# Attention types:\n",
        "RELATIVE16 = \"relative16\"\n",
        "RELATIVE16_BIASONLY = \"relative16_biasonly\"\n",
        "RELATIVE16_BIAS = \"relative16_bias\"\n",
        "\n",
        "RELATIVE16_D2E = \"relative16_d2e\"  # also relative positions for dec2enc\n",
        "RELATIVE16_D2E_BIASONLY = \"relative16_d2e_biasonly\"\n",
        "RELATIVE16_D2E_BIAS = \"relative16_d2e_bias\"\n",
        "\n",
        "ABSOLUTE_SINUSOIDAL = \"absolute_sinusoidal\"\n",
        "\n",
        "RELATIVE = \"relative\"\n",
        "RELATIVE_BIASONLY = \"relative_biasonly\"\n",
        "RELATIVE_BIAS = \"relative_bias\"\n",
        "\n",
        "# COPY_DECODER2 factors the copy decoder feedforward layer, to\n",
        "# prevent it from exploding in datasets with large vocabs:\n",
        "STANDARD_DECODER = \"standard_decoder\"\n",
        "COPY_DECODER = \"copy_decoder\"\n",
        "COPY_DECODER2 = \"copy_decoder2\"\n",
        "\n",
        "NO_REALFORMER = \"no_realformer\"\n",
        "REALFORMER_LOGITS = \"realformer_logits\"\n",
        "REALFORMER_LOGITSPSF = \"realformer_logits_plus_softmax\"\n",
        "\n",
        "STANDARD_FF = \"standard_ff\"\n",
        "\n",
        "GREEDY_DECODING = \"greedy_decoding\"\n",
        "\n",
        "\n",
        "def get_angles(pos, i, d_model):\n",
        "  angle_rates = 1 / np.power(10000, (2 * (i//2)) / np.float32(d_model))\n",
        "  return pos * angle_rates\n",
        "\n",
        "def positional_encoding(position, d_model):\n",
        "  angle_rads = get_angles(np.arange(position)[:, np.newaxis],\n",
        "                          np.arange(d_model)[np.newaxis, :],\n",
        "                          d_model)\n",
        "\n",
        "  # apply sin to even indices in the array; 2i\n",
        "  angle_rads[:, 0::2] = np.sin(angle_rads[:, 0::2])\n",
        "\n",
        "  # apply cos to odd indices in the array; 2i+1\n",
        "  angle_rads[:, 1::2] = np.cos(angle_rads[:, 1::2])\n",
        "\n",
        "  pos_encoding = angle_rads[np.newaxis, ...]\n",
        "\n",
        "  return tf.cast(pos_encoding, dtype=tf.float32)\n",
        "\n",
        "def create_padding_mask(seq):\n",
        "  seq = tf.cast(tf.math.equal(seq, 0), tf.float32)\n",
        "\n",
        "  # add extra dimensions to add the padding\n",
        "  # to the attention logits.\n",
        "  return seq[:, tf.newaxis, tf.newaxis, :]  # (batch_size, 1, 1, seq_len)\n",
        "\n",
        "def create_look_ahead_mask(size):\n",
        "  mask = 1 - tf.linalg.band_part(tf.ones((size, size)), -1, 0)\n",
        "  return mask  # (seq_len, seq_len)\n",
        "\n",
        "\n",
        "def scaled_dot_product_relative_attention(q, k, v,\n",
        "                                          mask=None,\n",
        "                                          relative_ids=None,\n",
        "                                          relative_embeddings=None,\n",
        "                                          relative_biases=None,\n",
        "                                          attention_logits_residual=None):\n",
        "  \"\"\"Calculate the attention weights.\n",
        "  - If relative_ids or relative_embeddings is None, then this is equivalent to\n",
        "    regular scaled_dot_product_attention.\n",
        "  - q, k, v must have matching leading dimensions.\n",
        "  - k, v must have matching penultimate dimension, i.e.: seq_len_k = seq_len_v.\n",
        "  - The mask has different shapes depending on its type(padding or look ahead)\n",
        "    but it must be broadcastable for addition.\n",
        "  - relative_ids is an optional [seq_len_q, seq_len_k] with the relative ids.\n",
        "  - relative_embeddings is an optional dense layer that converts OHE ids to\n",
        "    embeddings.\n",
        "  - relative_biases: optional dense layer that generates a bias term for the\n",
        "    attention logits based on the relative_ids\n",
        "  - attention_logits_residual is an optional parameter to implement\n",
        "    Realformer-style residual connections between Transformer layers.\n",
        "\n",
        "  Args:\n",
        "    q: query shape == (..., seq_len_q, depth)\n",
        "    k: key shape == (..., seq_len_k, depth)\n",
        "    v: value shape == (..., seq_len_v, depth_v)\n",
        "    mask: Float tensor with shape broadcastable\n",
        "          to (..., seq_len_q, seq_len_k). Defaults to None.\n",
        "    relative_ids == (seq_len_q, seq_len_k). Defaults to None.\n",
        "    relative_embeddings == dense layer. Defaults to None.\n",
        "\n",
        "  Returns:\n",
        "    output, attention_weights\n",
        "  \"\"\"\n",
        "\n",
        "  matmul_qk = tf.matmul(q, k, transpose_b=True)  # (..., seq_len_q, seq_len_k)\n",
        "\n",
        "  if relative_ids is not None:\n",
        "    if relative_embeddings is not None:\n",
        "      r = relative_embeddings(relative_ids)\n",
        "      matmul_qrel = tf.einsum(\"bhqd,qkd-\u003ebhqk\", q, r)\n",
        "      matmul_qk += matmul_qrel\n",
        "    if relative_biases is not None:\n",
        "      matmul_qk += tf.squeeze(relative_biases(relative_ids), axis=-1)\n",
        "\n",
        "  # scale matmul_qk\n",
        "  dk = tf.cast(tf.shape(k)[-1], tf.float32)\n",
        "  scaled_attention_logits = matmul_qk / tf.math.sqrt(dk)\n",
        "\n",
        "  # add the mask to the scaled tensor.\n",
        "  if mask is not None:\n",
        "    scaled_attention_logits += (mask * -1e9)\n",
        "\n",
        "  # Realformer style resitual connection:\n",
        "  if attention_logits_residual is not None:\n",
        "    scaled_attention_logits +=  attention_logits_residual\n",
        "\n",
        "  # softmax is normalized on the last axis (seq_len_k) so that the scores\n",
        "  # add up to 1.\n",
        "  # (..., seq_len_q, seq_len_k)\n",
        "  attention_weights = tf.nn.softmax(scaled_attention_logits, axis=-1)\n",
        "\n",
        "  output = tf.matmul(attention_weights, v)  # (..., seq_len_q, depth_v)\n",
        "\n",
        "  return output, scaled_attention_logits\n",
        "\n",
        "\n",
        "class MultiHeadAttention(tf.keras.layers.Layer):\n",
        "  def __init__(self, d_model, num_heads, relative_radius,\n",
        "               position_encodings=RELATIVE):\n",
        "    super(MultiHeadAttention, self).__init__()\n",
        "    self.num_heads = num_heads\n",
        "    self.d_model = d_model\n",
        "    self.relative_vocab_size = relative_radius*2+1\n",
        "\n",
        "    assert d_model % self.num_heads == 0\n",
        "\n",
        "    self.depth = d_model // self.num_heads\n",
        "\n",
        "    self.wq = tf.keras.layers.Dense(d_model)\n",
        "    self.wk = tf.keras.layers.Dense(d_model)\n",
        "    self.wv = tf.keras.layers.Dense(d_model)\n",
        "\n",
        "    self.dense = tf.keras.layers.Dense(d_model)\n",
        "\n",
        "    if (position_encodings == RELATIVE or\n",
        "        position_encodings == RELATIVE_BIAS):\n",
        "      self.relative_embeddings = tf.keras.layers.Embedding(\n",
        "          self.relative_vocab_size, self.depth)\n",
        "    else:\n",
        "      self.relative_embeddings = None\n",
        "    if (position_encodings == RELATIVE_BIAS or\n",
        "        position_encodings == RELATIVE_BIASONLY):\n",
        "      self.relative_biases = tf.keras.layers.Embedding(\n",
        "          self.relative_vocab_size, 1)\n",
        "    else:\n",
        "      self.relative_biases = None\n",
        "\n",
        "  def split_heads(self, x, batch_size):\n",
        "    \"\"\"Split the last dimension into (num_heads, depth).\n",
        "    Transpose the result such that the shape is (batch_size, num_heads, seq_len,\n",
        "    depth)\n",
        "    \"\"\"\n",
        "    x = tf.reshape(x, (batch_size, -1, self.num_heads, self.depth))\n",
        "    return tf.transpose(x, perm=[0, 2, 1, 3])\n",
        "\n",
        "  def call(self, v, k, q,\n",
        "           mask=None,\n",
        "           relative_ids=None,\n",
        "           attention_logits_residual=None):\n",
        "    batch_size = tf.shape(q)[0]\n",
        "\n",
        "    q = self.wq(q)  # (batch_size, seq_len, d_model)\n",
        "    k = self.wk(k)  # (batch_size, seq_len, d_model)\n",
        "    v = self.wv(v)  # (batch_size, seq_len, d_model)\n",
        "\n",
        "    q = self.split_heads(q, batch_size)  # (batch_size, num_heads, seq_len_q, depth)\n",
        "    k = self.split_heads(k, batch_size)  # (batch_size, num_heads, seq_len_k, depth)\n",
        "    v = self.split_heads(v, batch_size)  # (batch_size, num_heads, seq_len_v, depth)\n",
        "\n",
        "    # scaled_attention.shape == (batch_size, num_heads, seq_len_q, depth)\n",
        "    # attention_weights.shape == (batch_size, num_heads, seq_len_q, seq_len_k)\n",
        "    scaled_attention, attention_logits = scaled_dot_product_relative_attention(\n",
        "        q, k, v, mask, relative_ids, self.relative_embeddings,\n",
        "        self.relative_biases,\n",
        "        attention_logits_residual)\n",
        "\n",
        "    # (batch_size, seq_len_q, num_heads, depth)\n",
        "    scaled_attention = tf.transpose(scaled_attention, perm=[0, 2, 1, 3])\n",
        "\n",
        "    # (batch_size, seq_len_q, d_model)\n",
        "    concat_attention = tf.reshape(scaled_attention,\n",
        "                                  (batch_size, -1, self.d_model))\n",
        "\n",
        "    # (batch_size, seq_len_q, d_model)\n",
        "    output = self.dense(concat_attention)\n",
        "\n",
        "    return output, attention_logits\n",
        "\n",
        "\n",
        "def point_wise_feed_forward_network(d_model, dff):\n",
        "  return tf.keras.Sequential([\n",
        "      tf.keras.layers.Dense(dff,\n",
        "                            activation='relu'),  # (batch_size, seq_len, dff)\n",
        "      tf.keras.layers.Dense(d_model)  # (batch_size, seq_len, d_model)\n",
        "  ])\n",
        "\n",
        "\n",
        "class EncoderLayer(tf.keras.layers.Layer):\n",
        "  def __init__(self, d_model, num_heads, dff, rate=0.1, relative_radius=8,\n",
        "               position_encodings=RELATIVE):\n",
        "    super(EncoderLayer, self).__init__()\n",
        "\n",
        "    self.mha = MultiHeadAttention(d_model, num_heads, relative_radius,\n",
        "                                  position_encodings)\n",
        "    self.ffn = point_wise_feed_forward_network(d_model, dff)\n",
        "\n",
        "    self.layernorm1 = tf.keras.layers.LayerNormalization(epsilon=1e-6)\n",
        "    self.layernorm2 = tf.keras.layers.LayerNormalization(epsilon=1e-6)\n",
        "\n",
        "    self.dropout1 = tf.keras.layers.Dropout(rate)\n",
        "    self.dropout2 = tf.keras.layers.Dropout(rate)\n",
        "\n",
        "  def call(self, x, training, mask, relative_ids,\n",
        "           attention_logits_residual=None):\n",
        "\n",
        "    attn_output, attention_logits = self.mha(\n",
        "        x, x, x, mask, relative_ids,\n",
        "        attention_logits_residual)  # (batch_size, input_seq_len, d_model)\n",
        "    attn_output = self.dropout1(attn_output, training=training)\n",
        "    out1 = self.layernorm1(\n",
        "        x + attn_output)  # (batch_size, input_seq_len, d_model)\n",
        "\n",
        "    ffn_output = self.ffn(out1)  # (batch_size, input_seq_len, d_model)\n",
        "    ffn_output = self.dropout2(ffn_output, training=training)\n",
        "    out2 = self.layernorm2(\n",
        "        out1 + ffn_output)  # (batch_size, input_seq_len, d_model)\n",
        "\n",
        "    return out2, attention_logits\n",
        "\n",
        "class DecoderLayer(tf.keras.layers.Layer):\n",
        "  def __init__(self, d_model, num_heads, dff, rate=0.1, relative_radius=8,\n",
        "               position_encodings=RELATIVE):\n",
        "    super(DecoderLayer, self).__init__()\n",
        "\n",
        "    self.mha1 = MultiHeadAttention(d_model, num_heads, relative_radius,\n",
        "                                   position_encodings)\n",
        "    self.mha2 = MultiHeadAttention(d_model, num_heads, relative_radius,\n",
        "                                   position_encodings)\n",
        "\n",
        "    self.ffn = point_wise_feed_forward_network(d_model, dff)\n",
        "\n",
        "    self.layernorm1 = tf.keras.layers.LayerNormalization(epsilon=1e-6)\n",
        "    self.layernorm2 = tf.keras.layers.LayerNormalization(epsilon=1e-6)\n",
        "    self.layernorm3 = tf.keras.layers.LayerNormalization(epsilon=1e-6)\n",
        "\n",
        "    self.dropout1 = tf.keras.layers.Dropout(rate)\n",
        "    self.dropout2 = tf.keras.layers.Dropout(rate)\n",
        "    self.dropout3 = tf.keras.layers.Dropout(rate)\n",
        "\n",
        "\n",
        "  def call(self, x, enc_output, training,\n",
        "           look_ahead_mask, padding_mask,\n",
        "           dec_relative_ids, dec2enc_relative_ids,\n",
        "           attention_logits_residual1,\n",
        "           attention_logits_residual2):\n",
        "    # enc_output.shape == (batch_size, input_seq_len, d_model)\n",
        "    attn1, attn_logits_block1 = self.mha1(\n",
        "        x, x, x, look_ahead_mask, dec_relative_ids,\n",
        "        attention_logits_residual1)  # (batch_size, target_seq_len, d_model)\n",
        "    attn1 = self.dropout1(attn1, training=training)\n",
        "    out1 = self.layernorm1(attn1 + x)\n",
        "\n",
        "    attn2, attn_logits_block2 = self.mha2(\n",
        "        enc_output, enc_output, out1,\n",
        "        padding_mask, dec2enc_relative_ids,\n",
        "        attention_logits_residual2)  # (batch_size, target_seq_len, d_model)\n",
        "    attn2 = self.dropout2(attn2, training=training)\n",
        "    out2 = self.layernorm2(\n",
        "        attn2 + out1)  # (batch_size, target_seq_len, d_model)\n",
        "\n",
        "    ffn_output = self.ffn(out2)  # (batch_size, target_seq_len, d_model)\n",
        "    ffn_output = self.dropout3(ffn_output, training=training)\n",
        "    out3 = self.layernorm3(\n",
        "        ffn_output + out2)  # (batch_size, target_seq_len, d_model)\n",
        "\n",
        "    return out3, attn_logits_block1, attn_logits_block2\n",
        "\n",
        "\n",
        "class Encoder(tf.keras.layers.Layer):\n",
        "  def __init__(self, num_layers, d_model, num_heads, dff, input_vocab_size,\n",
        "               maximum_position_encoding, rate=0.1, relative_radius=8,\n",
        "               position_encodings=RELATIVE,\n",
        "               shared_layers=False,\n",
        "               realformer_connections=False):\n",
        "    super(Encoder, self).__init__()\n",
        "\n",
        "    self.d_model = d_model\n",
        "    self.num_layers = num_layers\n",
        "    self.position_encodings = position_encodings\n",
        "    self.shared_layers = shared_layers\n",
        "    self.realformer_connections = realformer_connections\n",
        "\n",
        "    self.embedding = tf.keras.layers.Embedding(input_vocab_size, d_model)\n",
        "    if self.position_encodings == ABSOLUTE_SINUSOIDAL:\n",
        "      self.pos_encoding = positional_encoding(maximum_position_encoding,\n",
        "                                              self.d_model)\n",
        "\n",
        "\n",
        "    if self.shared_layers:\n",
        "      layer = EncoderLayer(d_model, num_heads, dff, rate,\n",
        "                           relative_radius=relative_radius,\n",
        "                           position_encodings=position_encodings)\n",
        "      self.enc_layers = [layer for _ in range(num_layers)]\n",
        "    else:\n",
        "      self.enc_layers = [EncoderLayer(d_model, num_heads, dff, rate,\n",
        "                                      relative_radius=relative_radius,\n",
        "                                      position_encodings=position_encodings)\n",
        "                         for _ in range(num_layers)]\n",
        "\n",
        "    self.dropout = tf.keras.layers.Dropout(rate)\n",
        "\n",
        "  def call(self, x, training, mask, relative_ids):\n",
        "\n",
        "    seq_len = tf.shape(x)[1]\n",
        "\n",
        "    # adding embedding and position encoding.\n",
        "    x = self.embedding(x)  # (batch_size, input_seq_len, d_model)\n",
        "    x *= tf.math.sqrt(tf.cast(self.d_model, tf.float32))\n",
        "    if self.position_encodings == ABSOLUTE_SINUSOIDAL:\n",
        "      relative_ids = None\n",
        "      x += self.pos_encoding[:, :seq_len, :]\n",
        "\n",
        "    x = self.dropout(x, training=training)\n",
        "\n",
        "    attention_for_viz = []  # we accumulate attention logits for visualization\n",
        "    attention_logits_residual = None\n",
        "    for i in range(self.num_layers):\n",
        "      x, attention_logits_residual = self.enc_layers[i](\n",
        "          x, training, mask=mask, relative_ids=relative_ids,\n",
        "          attention_logits_residual=attention_logits_residual)\n",
        "      attention_for_viz.append(attention_logits_residual)\n",
        "      if self.realformer_connections == NO_REALFORMER:\n",
        "        attention_logits_residual = None\n",
        "\n",
        "    return x, attention_for_viz  # x: (batch_size, input_seq_len, d_model)\n",
        "\n",
        "  def detailed_param_count(self):\n",
        "    print(f\"  encoder embedding: {self.embedding.count_params()}\")\n",
        "    if self.shared_layers:\n",
        "      print(f\"  encoder layer weights (shared): {self.enc_layers[0].count_params()}\")\n",
        "    else:\n",
        "      print(f\"  encoder layer weights: {self.enc_layers[0].count_params()*len(self.enc_layers)}\")\n",
        "\n",
        "\n",
        "class Decoder(tf.keras.layers.Layer):\n",
        "  def __init__(self, num_layers, d_model, num_heads, dff, target_vocab_size,\n",
        "               maximum_position_encoding, rate=0.1, relative_radius=8,\n",
        "               position_encodings=RELATIVE,\n",
        "               shared_layers=False,\n",
        "               realformer_connections=NO_REALFORMER):\n",
        "    super(Decoder, self).__init__()\n",
        "\n",
        "    self.d_model = d_model\n",
        "    self.num_layers = num_layers\n",
        "    self.position_encodings = position_encodings\n",
        "    self.shared_layers = shared_layers\n",
        "    self.realformer_connections = realformer_connections\n",
        "\n",
        "    self.embedding = tf.keras.layers.Embedding(target_vocab_size, d_model)\n",
        "    if self.position_encodings == ABSOLUTE_SINUSOIDAL:\n",
        "      self.pos_encoding = positional_encoding(maximum_position_encoding,\n",
        "                                              d_model)\n",
        "\n",
        "    if self.shared_layers:\n",
        "      layer = DecoderLayer(d_model, num_heads, dff, rate,\n",
        "                           relative_radius=relative_radius,\n",
        "                           position_encodings=position_encodings)\n",
        "      self.dec_layers = [layer for _ in range(num_layers)]\n",
        "    else:\n",
        "      self.dec_layers = [DecoderLayer(d_model, num_heads, dff, rate,\n",
        "                                      relative_radius=relative_radius,\n",
        "                                      position_encodings=position_encodings)\n",
        "                        for _ in range(num_layers)]\n",
        "    self.dropout = tf.keras.layers.Dropout(rate)\n",
        "\n",
        "  def call(self, x, enc_output, training,\n",
        "           look_ahead_mask, padding_mask,\n",
        "           dec_relative_ids, dec2enc_relative_ids):\n",
        "\n",
        "    seq_len = tf.shape(x)[1]\n",
        "\n",
        "    x = self.embedding(x)  # (batch_size, target_seq_len, d_model)\n",
        "    x *= tf.math.sqrt(tf.cast(self.d_model, tf.float32))\n",
        "    if self.position_encodings == ABSOLUTE_SINUSOIDAL:\n",
        "      dec_relative_ids = None\n",
        "      dec2enc_relative_ids = None\n",
        "      x += self.pos_encoding[:, :seq_len, :]\n",
        "\n",
        "    x = self.dropout(x, training=training)\n",
        "\n",
        "    dec2dec_attention_for_viz = []\n",
        "    dec2enc_attention_for_viz = []\n",
        "    block1 = None\n",
        "    block2 = None\n",
        "    for i in range(self.num_layers):\n",
        "      x, block1, block2 = self.dec_layers[i](\n",
        "          x, enc_output, training, look_ahead_mask, padding_mask,\n",
        "          dec_relative_ids, dec2enc_relative_ids,\n",
        "          attention_logits_residual1 = block1,\n",
        "          attention_logits_residual2 = block2)\n",
        "      dec2dec_attention_for_viz.append(block1)\n",
        "      dec2enc_attention_for_viz.append(block2)\n",
        "      if self.realformer_connections == NO_REALFORMER:\n",
        "        block1 = None\n",
        "        block2 = None\n",
        "\n",
        "    # x.shape == (batch_size, target_seq_len, d_model)\n",
        "    return x, dec2dec_attention_for_viz, dec2enc_attention_for_viz\n",
        "\n",
        "  def detailed_param_count(self):\n",
        "    print(f\"  decoder embedding: {self.embedding.count_params()}\")\n",
        "    if self.shared_layers:\n",
        "      print(f\"  decoder layer weights (shared): {self.dec_layers[0].count_params()}\")\n",
        "    else:\n",
        "      print(f\"  decoder layer weights: {self.dec_layers[0].count_params()*len(self.dec_layers)}\")\n",
        "\n",
        "\n",
        "class Transformer(tf.keras.Model):\n",
        "  def __init__(self, num_layers, d_model, num_heads, dff, vocab_size,\n",
        "               pe_input, pe_target, rate=0.1,\n",
        "               relative_radius=8, position_encodings=RELATIVE,\n",
        "               copy_decoder=STANDARD_DECODER,\n",
        "               shared_layers=False,\n",
        "               realformer_connections=False):\n",
        "    super(Transformer, self).__init__()\n",
        "    self.d_model = d_model\n",
        "    self.vocab_size = vocab_size\n",
        "    self.copy_decoder = copy_decoder\n",
        "\n",
        "    self.encoder = Encoder(num_layers, d_model, num_heads, dff,\n",
        "                           vocab_size, pe_input, rate,\n",
        "                           relative_radius=relative_radius,\n",
        "                           position_encodings=position_encodings,\n",
        "                           shared_layers=shared_layers,\n",
        "                           realformer_connections=realformer_connections)\n",
        "\n",
        "    self.decoder = Decoder(num_layers, d_model, num_heads, dff,\n",
        "                           vocab_size, pe_target, rate,\n",
        "                           relative_radius=relative_radius,\n",
        "                           position_encodings=position_encodings,\n",
        "                           shared_layers=shared_layers,\n",
        "                           realformer_connections=realformer_connections)\n",
        "\n",
        "    self.final_layer = tf.keras.layers.Dense(vocab_size)\n",
        "    if self.copy_decoder == COPY_DECODER:\n",
        "      # pe_input is the maximum input length\n",
        "      self.final_layer_copy = tf.keras.layers.Dense(d_model)\n",
        "      self.final_layer_copy_weight = tf.keras.layers.Dense(\n",
        "          1, activation=\"sigmoid\")\n",
        "    elif self.copy_decoder == COPY_DECODER2:\n",
        "      self.final_layer_copy = tf.keras.layers.Dense(d_model)\n",
        "      self.final_layer_copy_weight = tf.keras.layers.Dense(\n",
        "          1, activation=\"sigmoid\")\n",
        "      # We want vocab_size -\u003e vocab_size, but that might be too big.\n",
        "      # So, we do a low-rank approximation, bringing it down to d_model first,\n",
        "      # in case d_model \u003c vocab_size:\n",
        "      if self.d_model \u003c self.vocab_size:\n",
        "        self.final_layer_copy2a = tf.keras.layers.Dense(self.d_model)\n",
        "        self.final_layer_copy2b = tf.keras.layers.Dense(self.vocab_size)\n",
        "      else:\n",
        "        self.final_layer_copy2 = tf.keras.layers.Dense(self.vocab_size)\n",
        "\n",
        "\n",
        "  def call(self, inp, tar, training, enc_padding_mask,\n",
        "           look_ahead_mask, dec_padding_mask,\n",
        "           enc_relative_ids, dec_relative_ids, dec2enc_relative_ids):\n",
        "\n",
        "    enc_output, enc2enc_attention_for_viz = (\n",
        "        self.encoder(inp, training, enc_padding_mask,\n",
        "                     enc_relative_ids))  # (batch_size, inp_seq_len, d_model)\n",
        "\n",
        "    # dec_output.shape == (batch_size, tar_seq_len, d_model)\n",
        "    (dec_output, dec2enc_attention_for_viz, dec2dec_attention_for_viz) = (\n",
        "        self.decoder(tar, enc_output, training, look_ahead_mask,\n",
        "                     dec_padding_mask,\n",
        "                     dec_relative_ids, dec2enc_relative_ids))\n",
        "\n",
        "    # (batch_size, tar_seq_len, vocab_size)\n",
        "    final_output = self.final_layer(dec_output)\n",
        "\n",
        "    # Copy decoder:\n",
        "    if (self.copy_decoder == COPY_DECODER or\n",
        "        self.copy_decoder == COPY_DECODER2):\n",
        "      # (batch_size, tar_seq_len, d_model)\n",
        "      copy_output_query = self.final_layer_copy(dec_output)\n",
        "      copy_output_weight = self.final_layer_copy_weight(dec_output)\n",
        "      # (batch_size, inp_seq_len, vocab_size)\n",
        "      copy_output, _ = scaled_dot_product_relative_attention(\n",
        "          copy_output_query,  # (batch_size, tar_seq_len, d_model)\n",
        "          enc_output,  # (batch_size, inp_seq_len, d_model)\n",
        "          tf.one_hot(inp, self.vocab_size))\n",
        "      if self.copy_decoder == COPY_DECODER2:\n",
        "        if self.d_model \u003c self.vocab_size:\n",
        "          copy_output = tf.nn.softmax(\n",
        "              self.final_layer_copy2b(\n",
        "                  self.final_layer_copy2a(copy_output)))\n",
        "        else:\n",
        "          copy_output = tf.nn.softmax(\n",
        "              self.final_layer_copy2(copy_output))\n",
        "      final_output = (\n",
        "          (1 - copy_output_weight) * tf.nn.softmax(final_output, axis=-1) +\n",
        "          copy_output_weight*copy_output)\n",
        "    else:\n",
        "      final_output = tf.nn.softmax(final_output, axis=-1)\n",
        "\n",
        "    return (final_output, enc2enc_attention_for_viz,\n",
        "            dec2enc_attention_for_viz, dec2dec_attention_for_viz)\n",
        "\n",
        "  def detailed_param_count(self):\n",
        "    print(\"Transformer parameter counts:\")\n",
        "    self.encoder.detailed_param_count()\n",
        "    self.decoder.detailed_param_count()\n",
        "    print(f\"  final_layer: {self.final_layer.count_params()}\")\n",
        "    if (self.copy_decoder == COPY_DECODER or\n",
        "        self.copy_decoder == COPY_DECODER2):\n",
        "      print(f\"  final_layer_copy: {self.final_layer_copy.count_params()}\")\n",
        "      print(f\"  final_layer_copy_weight: {self.final_layer_copy_weight.count_params()}\")\n",
        "    if self.copy_decoder == COPY_DECODER2:\n",
        "      if self.d_model \u003c self.vocab_size:\n",
        "        print(f\"  final_layer_copy2a: {self.final_layer_copy2a.count_params()}\")\n",
        "        print(f\"  final_layer_copy2b: {self.final_layer_copy2b.count_params()}\")\n",
        "      else:\n",
        "        print(f\"  final_layer_copy2: {self.final_layer_copy2.count_params()}\")\n",
        "\n",
        "class CustomSchedule(tf.keras.optimizers.schedules.LearningRateSchedule):\n",
        "  def __init__(self, d_model, warmup_steps=4000):\n",
        "    super(CustomSchedule, self).__init__()\n",
        "\n",
        "    self.d_model = d_model\n",
        "    self.d_model = tf.cast(self.d_model, tf.float32)\n",
        "\n",
        "    self.warmup_steps = warmup_steps\n",
        "\n",
        "  def __call__(self, step):\n",
        "    arg1 = tf.math.rsqrt(step)\n",
        "    arg2 = step * (self.warmup_steps ** -1.5)\n",
        "\n",
        "    return tf.math.rsqrt(self.d_model) * tf.math.minimum(arg1, arg2)\n",
        "\n",
        "\n",
        "def create_masks(inp, tar):\n",
        "  # Encoder padding mask\n",
        "  enc_padding_mask = create_padding_mask(inp)\n",
        "\n",
        "  # Used in the 2nd attention block in the decoder.\n",
        "  # This padding mask is used to mask the encoder outputs.\n",
        "  dec_padding_mask = create_padding_mask(inp)\n",
        "\n",
        "  # Used in the 1st attention block in the decoder.\n",
        "  # It is used to pad and mask future tokens in the input received by\n",
        "  # the decoder.\n",
        "  look_ahead_mask = create_look_ahead_mask(tf.shape(tar)[1])\n",
        "  dec_target_padding_mask = create_padding_mask(tar)\n",
        "  combined_mask = tf.maximum(dec_target_padding_mask, look_ahead_mask)\n",
        "\n",
        "  return enc_padding_mask, combined_mask, dec_padding_mask\n",
        "\n",
        "\n",
        "def create_relative_ids(inp_len, tar_len, relative_radius, de2enc_ids):\n",
        "  enc_relative_ids = np.zeros([inp_len, inp_len], dtype=int)\n",
        "  for i in range(inp_len):\n",
        "    for j in range(inp_len):\n",
        "      diff = i - j\n",
        "      diff = relative_radius + min(max(diff, -relative_radius), relative_radius)\n",
        "      enc_relative_ids[i][j] = diff\n",
        "\n",
        "  dec_relative_ids1 = np.zeros([tar_len-1, tar_len-1], dtype=int)\n",
        "  for i in range(tar_len-1):\n",
        "    for j in range(tar_len-1):\n",
        "      diff = i - j\n",
        "      diff = relative_radius + min(max(diff, -relative_radius), relative_radius)\n",
        "      dec_relative_ids1[i][j] = diff\n",
        "\n",
        "  dec2end_relative_ids = np.zeros([tar_len-1, inp_len], dtype=int)\n",
        "  for i in range(tar_len-1):\n",
        "    for j in range(inp_len):\n",
        "      if de2enc_ids:\n",
        "        diff = i - j\n",
        "        diff = relative_radius + min(max(diff, -relative_radius),\n",
        "                                     relative_radius)\n",
        "        dec2end_relative_ids[i][j] = diff\n",
        "      else:\n",
        "        dec2end_relative_ids[i][j] = relative_radius\n",
        "\n",
        "  return (tf.constant(enc_relative_ids), tf.constant(dec_relative_ids1),\n",
        "          tf.constant(dec2end_relative_ids))"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "id": "lenjYYMafjDt"
      },
      "outputs": [],
      "source": [
        "#@title Experiment execution function\n",
        "\n",
        "def setup_dataset(dataset):\n",
        "\n",
        "  vocab = [PAD_TOKEN, SEP_TOKEN, END_TOKEN, START_TOKEN,\n",
        "          END_ITERATION_TOKEN]\n",
        "  vocab_to_int = {PAD_TOKEN:0, SEP_TOKEN:SEP_TOKEN_IDX,\n",
        "                  END_TOKEN:END_TOKEN_IDX,\n",
        "                  START_TOKEN:START_TOKEN_IDX,\n",
        "                  END_ITERATION_TOKEN:END_ITERATION_TOKEN_IDX}\n",
        "\n",
        "  if dataset == DATASET_ADDITION_ALIGNED:\n",
        "    return create_addition_dataset(200000, 512, vocab, vocab_to_int,\n",
        "                                reversedigits=False, leftpadding=12,\n",
        "                                addAlignmentTokens=False)\n",
        "  elif dataset == DATASET_ADDITION_NEGATIVES:\n",
        "    return create_addition_dataset(200000, 512, vocab, vocab_to_int,\n",
        "                                reversedigits=False, leftpadding=12,\n",
        "                                addAlignmentTokens=False,\n",
        "                                negativeProbability=0.25)\n",
        "  elif dataset == DATASET_SCAN_LENGTH:\n",
        "    return create_scan_length_dataset(vocab, vocab_to_int,\n",
        "                                    map_output_to_input=False)\n",
        "  elif dataset == DATASET_SCAN_ADD_JUMP:\n",
        "    return create_scan_add_jump_dataset(vocab, vocab_to_int,\n",
        "                                     map_output_to_input=False)\n",
        "  elif dataset == DATASET_PCFG_PRODUCTIVITY:\n",
        "    return create_pcfg_dataset(\"productivity\", vocab, vocab_to_int)\n",
        "  elif dataset == DATASET_PCFG_SYSTEMATICITY:\n",
        "    return create_pcfg_dataset(\"systematicity\", vocab, vocab_to_int)\n",
        "  elif dataset == DATASET_COGS_FULL:\n",
        "    return create_cogs_dataset(vocab, vocab_to_int, max_len=512)\n",
        "  elif dataset == DATASET_CFQ_MCD1:\n",
        "    return create_cfq_mcd1_dataset(vocab, vocab_to_int)\n",
        "  elif dataset == DATASET_CFQ_MCD1_INTERMEDIATE:\n",
        "    return create_cfq_mcd1_dataset(vocab, vocab_to_int,\n",
        "                                   simplify_cartesians=True)\n",
        "  elif dataset == DATASET_DUPLICATION:\n",
        "    return create_duplicating_dataset(200000, 1024, vocab, vocab_to_int)\n",
        "  elif dataset == DATASET_REVERSING:\n",
        "    return create_reversing_dataset(200000, 1024, vocab, vocab_to_int)\n",
        "  elif dataset == DATASET_CARTESIAN:\n",
        "    return create_cartesian_dataset(200000, 1024, vocab, vocab_to_int)\n",
        "  elif dataset == DATASET_INTERSECTION_BOOLEAN:\n",
        "    return create_intersection_dataset(200000, 1024, vocab, vocab_to_int)\n",
        "  else:\n",
        "    raise ValueError(f\"Undefined dataset type: {dataset}\")\n",
        "\n",
        "\n",
        "def run_evaluation(dataset,\n",
        "                   num_layers, d_model, dff, num_heads, position_encoding,\n",
        "                   decoder_type, share_layer_weights,\n",
        "                   num_epochs,\n",
        "                   realformer=NO_REALFORMER,\n",
        "                   num_repetitions=1,\n",
        "                   save_results_every_n_epochs=None,\n",
        "                   batch_size=64):\n",
        "\n",
        "  if save_results_every_n_epochs is None:\n",
        "    save_results_every_n_epochs = num_epochs\n",
        "\n",
        "  # Load and setup the corresponding dataset:\n",
        "  (vocab, vocab_to_int, input_tensor_train, target_tensor_train,\n",
        "    input_tensor_val_list, target_tensor_val_list) = setup_dataset(dataset)\n",
        "  (dataset_train, dataset_val_list) = prepare_tf_dataset_tensors(\n",
        "      vocab, vocab_to_int, input_tensor_train, target_tensor_train,\n",
        "      input_tensor_val_list, target_tensor_val_list, batch_size)\n",
        "  max_len_inp = len(input_tensor_train[0])\n",
        "  max_len_targ = len(target_tensor_train[0])\n",
        "\n",
        "  if position_encoding == RELATIVE16:\n",
        "    position_encoding = RELATIVE\n",
        "    relative_radius = 16\n",
        "    de2enc_ids = False\n",
        "  elif position_encoding == RELATIVE16_BIASONLY:\n",
        "    position_encoding = RELATIVE_BIASONLY\n",
        "    relative_radius = 16\n",
        "    de2enc_ids = False\n",
        "  elif position_encoding == RELATIVE16_BIAS:\n",
        "    position_encoding = RELATIVE_BIAS\n",
        "    relative_radius = 16\n",
        "    de2enc_ids = False\n",
        "  elif position_encoding == RELATIVE16_D2E:\n",
        "    position_encoding = RELATIVE\n",
        "    relative_radius = 16\n",
        "    de2enc_ids = True\n",
        "  elif position_encoding == RELATIVE16_D2E_BIASONLY:\n",
        "    position_encoding = RELATIVE_BIASONLY\n",
        "    relative_radius = 16\n",
        "    de2enc_ids = True\n",
        "  elif position_encoding == RELATIVE16_D2E_BIAS:\n",
        "    position_encoding = RELATIVE_BIAS\n",
        "    relative_radius = 16\n",
        "    de2enc_ids = True\n",
        "  elif position_encoding == ABSOLUTE_SINUSOIDAL:\n",
        "    position_encoding = ABSOLUTE_SINUSOIDAL\n",
        "    relative_radius = 0\n",
        "    de2enc_ids = False\n",
        "  else:\n",
        "    raise ValueError(f\"Undefined position embeddings type: {position_embeddings}\")\n",
        "\n",
        "  dropout_rate = 0.1\n",
        "\n",
        "  repetition_metrics_l = []\n",
        "  for repetition in range(num_repetitions):\n",
        "    print(f\"Starting repetition {repetition} ...\\n\")\n",
        "\n",
        "    learning_rate = CustomSchedule(d_model)\n",
        "    optimizer = tf.keras.optimizers.Adam(learning_rate, beta_1=0.9, beta_2=0.98,\n",
        "                                        epsilon=1e-9)\n",
        "    loss_object = tf.keras.losses.SparseCategoricalCrossentropy(\n",
        "        from_logits=False, reduction='none')\n",
        "    train_loss = tf.keras.metrics.Mean(name='train_loss')\n",
        "    train_accuracy = tf.keras.metrics.Mean(name='train_accuracy')\n",
        "    eval_loss = tf.keras.metrics.Mean(name='train_loss')\n",
        "    eval_accuracy = tf.keras.metrics.Mean(name='train_accuracy')\n",
        "\n",
        "    transformer = Transformer(num_layers, d_model, num_heads, dff,\n",
        "                              len(vocab),\n",
        "                              pe_input=max_len_inp,\n",
        "                              pe_target=max_len_targ,\n",
        "                              rate=dropout_rate,\n",
        "                              relative_radius=relative_radius,\n",
        "                              position_encodings=position_encoding,\n",
        "                              copy_decoder=decoder_type,\n",
        "                              shared_layers=share_layer_weights,\n",
        "                              realformer_connections=realformer)\n",
        "\n",
        "\n",
        "    def loss_function(real, pred):\n",
        "      mask = tf.math.logical_not(tf.math.equal(real, 0))\n",
        "      loss_ = loss_object(real, pred)\n",
        "\n",
        "      mask = tf.cast(mask, dtype=loss_.dtype)\n",
        "      loss_ *= mask\n",
        "\n",
        "      return tf.reduce_sum(loss_)/tf.reduce_sum(mask)\n",
        "\n",
        "\n",
        "    def accuracy_function(real, pred):\n",
        "      mask = tf.math.logical_not(tf.math.equal(real, 0))\n",
        "      loss_ = tf.keras.metrics.sparse_categorical_accuracy(real, pred)\n",
        "\n",
        "      mask = tf.cast(mask, dtype=loss_.dtype)\n",
        "      loss_ *= mask\n",
        "\n",
        "      return tf.reduce_sum(loss_)/tf.reduce_sum(mask)\n",
        "\n",
        "\n",
        "    train_step_signature = [\n",
        "        tf.TensorSpec(shape=(None, None), dtype=tf.int32),\n",
        "        tf.TensorSpec(shape=(None, None), dtype=tf.int32)\n",
        "    ]\n",
        "\n",
        "    @tf.function(input_signature=train_step_signature)\n",
        "    def train_step(inp, targ):\n",
        "      targ_inp = targ[:, :-1]\n",
        "      targ_real = targ[:, 1:]\n",
        "      enc_padding_mask, combined_mask, dec_padding_mask = create_masks(inp,\n",
        "                                                                       targ_inp)\n",
        "      enc_relative_ids, dec_relative_ids, dec2end_relative_ids = (\n",
        "          create_relative_ids(max_len_inp, max_len_targ, relative_radius,\n",
        "                              de2enc_ids))\n",
        "      with tf.GradientTape() as tape:\n",
        "        predictions, _, _, _ = transformer(inp, targ_inp,\n",
        "                                           True,\n",
        "                                           enc_padding_mask,\n",
        "                                           combined_mask,\n",
        "                                           dec_padding_mask,\n",
        "                                           enc_relative_ids,\n",
        "                                           dec_relative_ids,\n",
        "                                           dec2end_relative_ids)\n",
        "        loss = loss_function(targ_real, predictions)\n",
        "        accuracy = accuracy_function(targ_real, predictions)\n",
        "\n",
        "      gradients = tape.gradient(loss, transformer.trainable_variables)\n",
        "      optimizer.apply_gradients(zip(gradients, transformer.trainable_variables))\n",
        "      train_loss(loss)\n",
        "      train_accuracy(accuracy)\n",
        "\n",
        "\n",
        "    @tf.function(input_signature=train_step_signature)\n",
        "    def eval_step(inp, targ):\n",
        "      targ_inp = targ[:, :-1]\n",
        "      targ_real = targ[:, 1:]\n",
        "      enc_padding_mask, combined_mask, dec_padding_mask = create_masks(inp,\n",
        "                                                                       targ_inp)\n",
        "      enc_relative_ids, dec_relative_ids, dec2end_relative_ids = (\n",
        "          create_relative_ids(max_len_inp, max_len_targ, relative_radius,\n",
        "                              de2enc_ids))\n",
        "      with tf.GradientTape() as tape:\n",
        "        predictions, _, _, _ = transformer(inp, targ_inp,\n",
        "                                           False,\n",
        "                                           enc_padding_mask,\n",
        "                                           combined_mask,\n",
        "                                           dec_padding_mask,\n",
        "                                           enc_relative_ids,\n",
        "                                           dec_relative_ids,\n",
        "                                           dec2end_relative_ids)\n",
        "        loss = loss_function(targ_real, predictions)\n",
        "        accuracy = accuracy_function(targ_real, predictions)\n",
        "      eval_loss(loss)\n",
        "      eval_accuracy(accuracy)\n",
        "\n",
        "\n",
        "    # This method is slow, so, we just want to call it once per epoch, to\n",
        "    # visualize how the model is doing.\n",
        "    def eval_step_detailed(inp, targ, max_to_show, vocab):\n",
        "      targ_inp = targ[:, :-1]\n",
        "      targ_real = targ[:, 1:]\n",
        "      enc_padding_mask, combined_mask, dec_padding_mask = create_masks(inp,\n",
        "                                                                       targ_inp)\n",
        "      enc_relative_ids, dec_relative_ids, dec2end_relative_ids = (\n",
        "          create_relative_ids(max_len_inp, max_len_targ, relative_radius,\n",
        "                              de2enc_ids))\n",
        "      predictions, _, _, _ = transformer(inp, targ_inp,\n",
        "                                         False,\n",
        "                                         enc_padding_mask,\n",
        "                                         combined_mask,\n",
        "                                         dec_padding_mask,\n",
        "                                         enc_relative_ids,\n",
        "                                         dec_relative_ids,\n",
        "                                         dec2end_relative_ids)\n",
        "\n",
        "      predicted_targ = predicted_ids = tf.argmax(predictions, 2).numpy()\n",
        "\n",
        "      accuracy_seq = 0\n",
        "      accuracy_tok = 0\n",
        "      total_tok = 0\n",
        "      shown = 0\n",
        "      inp_numpy = inp.numpy()\n",
        "      ground_truth = targ_real.numpy()\n",
        "      for i in range(batch_size):\n",
        "        for j in range(int(targ_real.shape[1])-1):\n",
        "          if predicted_targ[i][j] == 2 or predicted_targ[i][j] == 0:\n",
        "            predicted_targ[i][j+1] = 0\n",
        "        for k in range(len(predicted_targ[i])):\n",
        "          if predicted_targ[i][k] == ground_truth[i][k]:\n",
        "            accuracy_tok += 1\n",
        "          total_tok += 1\n",
        "        if (predicted_targ[i] == ground_truth[i]).all():\n",
        "          accuracy_seq += 1\n",
        "        else:\n",
        "          if shown \u003c max_to_show:\n",
        "            print(\"input:     \" + decode(inp_numpy[i], vocab))\n",
        "            print(\"target:    \" + decode(ground_truth[i], vocab))\n",
        "            print(\"predicted: \" + decode(predicted_targ[i], vocab))\n",
        "            shown += 1\n",
        "      return accuracy_tok / float(total_tok), accuracy_seq / float(batch_size)\n",
        "\n",
        "    def evaluate_in_set(dataset_val, vocab):\n",
        "      steps_per_epoch_val = dataset_val.cardinality()\n",
        "      n_test_batches = 0\n",
        "      test_accuracy_token = 0\n",
        "      test_accuracy_seq = 0\n",
        "      for (batch, (inp, targ)) in enumerate(dataset_val.take(\n",
        "          steps_per_epoch_val)):\n",
        "        if batch == 0:\n",
        "          batch_accuracy_token, batch_accuracy_seq = eval_step_detailed(\n",
        "              inp, targ, 16, vocab)\n",
        "          test_accuracy_token += batch_accuracy_token\n",
        "          test_accuracy_seq += batch_accuracy_seq\n",
        "        else:\n",
        "          batch_accuracy_token, batch_accuracy_seq = eval_step_detailed(\n",
        "              inp, targ, 0, vocab)\n",
        "          test_accuracy_token += batch_accuracy_token\n",
        "          test_accuracy_seq += batch_accuracy_seq\n",
        "        n_test_batches += 1\n",
        "\n",
        "      print (f'Eval accuracy (token level): {test_accuracy_token/n_test_batches}')\n",
        "      print (f'Eval accuracy (sequence level): {test_accuracy_seq/n_test_batches}\\n')\n",
        "      return (test_accuracy_token/n_test_batches, test_accuracy_seq/n_test_batches)\n",
        "\n",
        "    steps_per_epoch = len(input_tensor_train)//batch_size\n",
        "    repetition_metrics = []\n",
        "    for epoch in range(num_epochs):\n",
        "      steps_per_epoch_val = dataset_val_list[0].cardinality()\n",
        "      start = time.time()\n",
        "\n",
        "      train_loss.reset_states()\n",
        "      train_accuracy.reset_states()\n",
        "      eval_loss.reset_states()\n",
        "      eval_accuracy.reset_states()\n",
        "\n",
        "      for (batch, (inp, targ)) in enumerate(dataset_train.take(steps_per_epoch)):\n",
        "        # print(\"inp:\" + str(inp.shape))\n",
        "        # print(\"targ:\" + str(targ.shape))\n",
        "        train_step(inp, targ)\n",
        "        if batch % 100 == 0:\n",
        "          print ('Epoch {} Batch {} Loss {:.4f} Accuracy {:.4f}'.format(\n",
        "              epoch + 1, batch, train_loss.result(), train_accuracy.result()))\n",
        "      print ('Epoch {} Loss {:.4f} Accuracy {:.4f}'.format(epoch + 1,\n",
        "                                                    train_loss.result(),\n",
        "                                                    train_accuracy.result()))\n",
        "      for (batch, (inp, targ)) in enumerate(dataset_val_list[0].take(\n",
        "          steps_per_epoch_val)):\n",
        "        eval_step(inp, targ)\n",
        "      print ('Epoch {} Eval Loss {:.4f} Eval Accuracy {:.4f}'.format(epoch + 1,\n",
        "                                                    eval_loss.result(),\n",
        "                                                    eval_accuracy.result()))\n",
        "\n",
        "      print ('Time taken for 1 epoch: {} secs\\n'.format(time.time() - start))\n",
        "\n",
        "      if ((epoch+1) % save_results_every_n_epochs) == 0:\n",
        "        epoch_metrics = [epoch+1]\n",
        "        for i in range(len(dataset_val_list)):\n",
        "          print(f\"------- Evaluation in dataset_val {i} -------\")\n",
        "          (acc_token, acc_seq) = evaluate_in_set(dataset_val_list[i], vocab)\n",
        "          epoch_metrics.append(acc_token)\n",
        "          epoch_metrics.append(acc_seq)\n",
        "        repetition_metrics.append(epoch_metrics)\n",
        "\n",
        "    # Make sure we save the last epoch results if it is not a\n",
        "    # multiple of save_results_every_n_epochs:\n",
        "    if (num_epochs % save_results_every_n_epochs) != 0:\n",
        "      epoch_metrics = [num_epochs]\n",
        "      for i in range(len(dataset_val_list)):\n",
        "        print(f\"------- Evaluation in dataset_val {i} -------\")\n",
        "        (acc_token, acc_seq) = evaluate_in_set(dataset_val_list[i], vocab)\n",
        "        epoch_metrics.append(acc_token)\n",
        "        epoch_metrics.append(acc_seq)\n",
        "      repetition_metrics.append(epoch_metrics)\n",
        "\n",
        "    repetition_metrics_l.append(repetition_metrics)\n",
        "\n",
        "  print(\"Raw repetition metrics:\")\n",
        "  for repetition_metrics in repetition_metrics_l:\n",
        "    for epoch_metrics in repetition_metrics:\n",
        "      print(epoch_metrics)\n",
        "\n",
        "  averages = repetition_metrics_l[0]\n",
        "  for i in range(1, len(repetition_metrics_l)):\n",
        "    for j in range(len(repetition_metrics_l[i])):\n",
        "      epoch_metrics = repetition_metrics_l[i][j]\n",
        "      # skip the epoch number:\n",
        "      for k in range(1, len(epoch_metrics)):\n",
        "        averages[j][k] += epoch_metrics[k]\n",
        "  for j in range(len(averages)):\n",
        "    for k in range(1, len(averages[j])):\n",
        "      averages[j][k] /= len(repetition_metrics_l)\n",
        "\n",
        "  print(\"Average repetition metrics:\")\n",
        "  for average_metrics in averages:\n",
        "    print('\\t'.join(map(str,average_metrics)))\n",
        "\n",
        "  print(f\"Transformer params: {transformer.count_params()}\")\n",
        "  transformer.detailed_param_count()\n",
        "\n",
        "  return averages, transformer, dataset_train, dataset_val_list\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "id": "1A8QBZM0pfnK"
      },
      "outputs": [],
      "source": [
        "#@title Individual Experiments\n",
        "\n",
        "# Run the different experiments:\n",
        "# parameters are:\n",
        "# - dataset\n",
        "# - num_layers, d_model, dff, num_heads\n",
        "# - position_encoding, decoder_type, share_layer_weights\n",
        "# - num_epochs\n",
        "# - optional keyword parameters:\n",
        "#     num_repetitions=1,\n",
        "#     save_results_every_n_epochs=None, batch_size=64\n",
        "\n",
        "# Example:\n",
        "run_evaluation(DATASET_DUPLICATION,\n",
        "               4, 128, 512, 8,\n",
        "               RELATIVE16_D2E_BIAS, COPY_DECODER, True,\n",
        "               4,\n",
        "               num_repetitions=1, save_results_every_n_epochs=1)\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "id": "njrypvBvNMQc"
      },
      "outputs": [],
      "source": [
        "#@title Systematic Experiments\n",
        "\n",
        "# Set the path where you want the results to be stored:\n",
        "RESULTS_FILE = \"/path-to/results.tsv\"\n",
        "\n",
        "datasets = [\n",
        "            (DATASET_ADDITION_ALIGNED, 2),\n",
        "            (DATASET_ADDITION_NEGATIVES, 10),\n",
        "            (DATASET_REVERSING, 2),\n",
        "            (DATASET_DUPLICATION, 4),\n",
        "            (DATASET_CARTESIAN, 4),\n",
        "            (DATASET_INTERSECTION_BOOLEAN, 8),\n",
        "            (DATASET_SCAN_LENGTH, 24),\n",
        "            (DATASET_SCAN_ADD_JUMP, 24),\n",
        "            (DATASET_PCFG_PRODUCTIVITY, 20),\n",
        "            (DATASET_PCFG_SYSTEMATICITY, 20),\n",
        "            (DATASET_COGS_FULL, 16),\n",
        "            (DATASET_CFQ_MCD1, 16),\n",
        "            (DATASET_CFQ_MCD1_INTERMEDIATE, 16)\n",
        "            ]\n",
        "\n",
        "num_repetitions = 5\n",
        "\n",
        "models = [\n",
        "          (\"abs\", 2, 64, 256, 4, ABSOLUTE_SINUSOIDAL, STANDARD_DECODER, False, NO_REALFORMER, STANDARD_FF, GREEDY_DECODING),\n",
        "          (\"rel-e\", 2, 64, 256, 4, RELATIVE16, STANDARD_DECODER, False, NO_REALFORMER, STANDARD_FF, GREEDY_DECODING),\n",
        "          (\"rel-b\", 2, 64, 256, 4, RELATIVE16_BIASONLY, STANDARD_DECODER, False, NO_REALFORMER, STANDARD_FF, GREEDY_DECODING),\n",
        "          (\"rel-eb\", 2, 64, 256, 4, RELATIVE16_BIAS, STANDARD_DECODER, False, NO_REALFORMER, STANDARD_FF, GREEDY_DECODING),\n",
        "          (\"rel2-e\", 2, 64, 256, 4, RELATIVE16_D2E, STANDARD_DECODER, False, NO_REALFORMER, STANDARD_FF, GREEDY_DECODING),\n",
        "          (\"rel2-b\", 2, 64, 256, 4, RELATIVE16_D2E_BIASONLY, STANDARD_DECODER, False, NO_REALFORMER, STANDARD_FF, GREEDY_DECODING),\n",
        "          (\"rel2-eb\", 2, 64, 256, 4, RELATIVE16_D2E_BIAS, STANDARD_DECODER, False, NO_REALFORMER, STANDARD_FF, GREEDY_DECODING),\n",
        "          (\"abs-c\", 2, 64, 256, 4, ABSOLUTE_SINUSOIDAL, COPY_DECODER, False, NO_REALFORMER, STANDARD_FF, GREEDY_DECODING),\n",
        "          (\"rel-eb-c\", 2, 64, 256, 4, RELATIVE16_BIAS, COPY_DECODER, False, NO_REALFORMER, STANDARD_FF, GREEDY_DECODING),\n",
        "          (\"rel2-eb-c\", 2, 64, 256, 4, RELATIVE16_D2E_BIAS, COPY_DECODER, False, NO_REALFORMER, STANDARD_FF, GREEDY_DECODING),\n",
        "          (\"small-2\", 2, 64, 256, 4, RELATIVE16_D2E_BIAS, COPY_DECODER, False, NO_REALFORMER, STANDARD_FF, GREEDY_DECODING),\n",
        "          (\"small-4\", 4, 64, 256, 4, RELATIVE16_D2E_BIAS, COPY_DECODER, False, NO_REALFORMER, STANDARD_FF, GREEDY_DECODING),\n",
        "          (\"small-6\", 6, 64, 256, 4, RELATIVE16_D2E_BIAS, COPY_DECODER, False, NO_REALFORMER, STANDARD_FF, GREEDY_DECODING),\n",
        "          (\"large-2\", 4, 128, 512, 8, RELATIVE16_D2E_BIAS, COPY_DECODER, False, NO_REALFORMER, STANDARD_FF, GREEDY_DECODING),\n",
        "          (\"large-4\", 4, 128, 512, 8, RELATIVE16_D2E_BIAS, COPY_DECODER, False, NO_REALFORMER, STANDARD_FF, GREEDY_DECODING),\n",
        "          (\"large-6\", 6, 128, 512, 8, RELATIVE16_D2E_BIAS, COPY_DECODER, False, NO_REALFORMER, STANDARD_FF, GREEDY_DECODING),\n",
        "          (\"small-2s\", 2, 64, 256, 4, RELATIVE16_D2E_BIAS, COPY_DECODER, True, NO_REALFORMER, STANDARD_FF, GREEDY_DECODING),\n",
        "          (\"small-4s\", 4, 64, 256, 4, RELATIVE16_D2E_BIAS, COPY_DECODER, True, NO_REALFORMER, STANDARD_FF, GREEDY_DECODING),\n",
        "          (\"small-6s\", 6, 64, 256, 4, RELATIVE16_D2E_BIAS, COPY_DECODER, True, NO_REALFORMER, STANDARD_FF, GREEDY_DECODING),\n",
        "          (\"large-2s\", 2, 128, 512, 8, RELATIVE16_D2E_BIAS, COPY_DECODER, True, NO_REALFORMER, STANDARD_FF, GREEDY_DECODING),\n",
        "          (\"large-4s\", 4, 128, 512, 8, RELATIVE16_D2E_BIAS, COPY_DECODER, True, NO_REALFORMER, STANDARD_FF, GREEDY_DECODING),\n",
        "          (\"large-6s\", 6, 128, 512, 8, RELATIVE16_D2E_BIAS, COPY_DECODER, True, NO_REALFORMER, STANDARD_FF, GREEDY_DECODING),\n",
        "          ]\n",
        "\n",
        "for (modelname, num_layers, d_model, dff, num_heads, position_encoding, decoder_type, share_layer_weights, realformer, ff_type, decoding_strategy) in models:\n",
        "  for (dataset, epochs) in datasets:\n",
        "    for i in range(num_repetitions):\n",
        "      results, transformer, _, _ = run_evaluation(\n",
        "          dataset,\n",
        "          num_layers, d_model, dff, num_heads, position_encoding, decoder_type, share_layer_weights, epochs, realformer=realformer,\n",
        "          save_results_every_n_epochs=2)\n",
        "      params = transformer.count_params()\n",
        "      with gfile.Open(RESULTS_FILE, \"a\") as wf:\n",
        "        for result in results:\n",
        "          report = [modelname, dataset, params] + result\n",
        "          report_str = \"\\t\".join([str(x) for x in report])\n",
        "          print(report_str)\n",
        "          wf.write(report_str)\n",
        "          wf.write(\"\\n\")\n",
        "\n"
      ]
    }
  ],
  "metadata": {
    "colab": {
      "collapsed_sections": [],
      "last_runtime": {
        "build_target": "//learning/deepmind/dm_python:dm_notebook3",
        "kind": "private"
      },
      "name": "seq2seq_tf.ipynb",
      "provenance": [
        {
          "file_id": "1yoke9fHsL6eHFq3MmVRoq82hpo7xHv1v",
          "timestamp": 1601564673821
        },
        {
          "file_id": "1I41hx_VfIBxmuXygVEwlZ277A5HHp6ze",
          "timestamp": 1600681566885
        }
      ]
    },
    "kernelspec": {
      "display_name": "Python 3",
      "name": "python3"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 0
}
